From dbbf845e56e77047b39e987e4c5762a212992f4e Mon Sep 17 00:00:00 2001
From: Sviatoslav Tsariov <slava982007@gmail.com>
Date: Wed, 20 Mar 2024 12:51:14 +0300
Subject: [PATCH] Added whisper and fast whisper implementation

---
 src/app.py                                    |  1 -
 src/message/message.py                        |  2 +-
 src/message/message_sender/message_sender.py  |  4 +--
 .../message_sender/message_sender_strategy.py |  2 +-
 src/recognizer/fast_whisper_strategy.py       | 34 ++++++++++++++++++-
 src/recognizer/whisper_strategy.py            |  9 ++++-
 6 files changed, 45 insertions(+), 7 deletions(-)

diff --git a/src/app.py b/src/app.py
index f42943d..25ac26a 100644
--- a/src/app.py
+++ b/src/app.py
@@ -32,7 +32,6 @@ def recognize_files(handler_fn):
     results = []
 
     for filename, handle in request.files.items():
-        # 
         temp = NamedTemporaryFile()
         handle.save(temp)
         result = model.transcribe(temp.name, language=HARPYIA_LANGUAGE, initial_prompt=HARPYIA_PROMPT)
diff --git a/src/message/message.py b/src/message/message.py
index d67274c..6eece7e 100644
--- a/src/message/message.py
+++ b/src/message/message.py
@@ -10,7 +10,7 @@ class Message:
   def has_prompt(self) -> bool:
     return self._prompt_service.has_prompt(self._recognized_text)
 
-  def send(self):
+  def send(self) -> None:
     self._message_sender.send(self._generate_response())
 
   def _generate_response(self) -> str:
diff --git a/src/message/message_sender/message_sender.py b/src/message/message_sender/message_sender.py
index d16e33b..4705e60 100644
--- a/src/message/message_sender/message_sender.py
+++ b/src/message/message_sender/message_sender.py
@@ -1,8 +1,8 @@
 from message.message_sender import MessageSenderStrategy
 
 class MessageSender:
-  def __init__(self, strategy: MessageSenderStrategy):
+  def __init__(self, strategy: MessageSenderStrategy) -> None:
     self._strategy = strategy
 
-  def send(self, message):
+  def send(self, message) -> None:
     self._strategy.send(message)
diff --git a/src/message/message_sender/message_sender_strategy.py b/src/message/message_sender/message_sender_strategy.py
index 1fc618a..7855a58 100644
--- a/src/message/message_sender/message_sender_strategy.py
+++ b/src/message/message_sender/message_sender_strategy.py
@@ -2,5 +2,5 @@ from abc import ABC, abstractmethod
 
 class MessageSenderStrategy(ABC):
   @abstractmethod
-  def send(self, message):
+  def send(self, message) -> None:
     pass
diff --git a/src/recognizer/fast_whisper_strategy.py b/src/recognizer/fast_whisper_strategy.py
index d93680d..458f932 100644
--- a/src/recognizer/fast_whisper_strategy.py
+++ b/src/recognizer/fast_whisper_strategy.py
@@ -1,5 +1,37 @@
+import whisper
+from faster_whisper import WhisperModel
+
+import config
 from recognizer import RecognizerStrategy
 
 class FastWhisperStrategy(RecognizerStrategy):
+  def __init__(self) -> None:
+      self._model = WhisperModel(
+      model_size=config.HARPYIA_MODEL,
+      device=config.DEVICE,
+      num_workers=6,
+      cpu_threads=10,
+      # in_memory=True,
+    )
+
   def recognize(self, file) -> str:
-    return ''
\ No newline at end of file
+      audio = self._prepare_file(file.name)
+      return self._transcribe(audio)
+
+  def _prepare_file(self, filename: str):
+    audio = whisper.load_audio(filename, sr=16000)
+    audio = whisper.pad_or_trim(audio)
+    return audio
+
+
+  def _transcribe(self, audio):
+    segments, _ = self._model.transcribe(
+      audio,
+      language=config.HARPYIA_LANGUAGE,
+      initial_prompt=config.HARPYIA_PROMPT,
+      condition_on_previous_text=False,
+      vad_filter=True,
+      beam_size=5,
+    )
+
+    return segments
\ No newline at end of file
diff --git a/src/recognizer/whisper_strategy.py b/src/recognizer/whisper_strategy.py
index 2fcf3af..2472085 100644
--- a/src/recognizer/whisper_strategy.py
+++ b/src/recognizer/whisper_strategy.py
@@ -1,5 +1,12 @@
+import whisper
+
+import config
 from recognizer import RecognizerStrategy
 
 class WhisperStrategy(RecognizerStrategy):
+  def __init__(self) -> None:
+    self._model = whisper.load_model(config.HARPYIA_MODEL, device=config.DEVICE)
+
   def recognize(self, file) -> str:
-    return ''
\ No newline at end of file
+    return self._model.transcribe(file.name, \
+      language=config.HARPYIA_LANGUAGE, initial_prompt=config.HARPYIA_PROMPT)
\ No newline at end of file