Adding support for Groq.com API (#813)

chidiwilliams · Jun 22, 2024 · cf340bc · cf340bc
1 parent 0046139
commit cf340bc
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 6 deletions.
diff --git a/buzz/model_loader.py b/buzz/model_loader.py
@@ -481,3 +481,10 @@ def download_model(
 
     def cancel(self):
         self.stopped = True
+
+
+def get_custom_api_whisper_model(base_url: str):
+    if "api.groq.com" in base_url:
+        return "whisper-large-v3"
+
+    return "whisper-1"
diff --git a/buzz/transcriber/openai_whisper_api_file_transcriber.py b/buzz/transcriber/openai_whisper_api_file_transcriber.py
@@ -9,6 +9,7 @@
 from openai import OpenAI
 
 from buzz.settings.settings import Settings
+from buzz.model_loader import get_custom_api_whisper_model
 from buzz.transcriber.file_transcriber import FileTranscriber
 from buzz.transcriber.transcriber import FileTranscriptionTask, Segment, Task
 
@@ -25,6 +26,9 @@ def __init__(self, task: FileTranscriptionTask, parent: Optional["QObject"] = No
             api_key=self.transcription_task.transcription_options.openai_access_token,
             base_url=custom_openai_base_url if custom_openai_base_url else None
         )
+        self.whisper_api_model = get_custom_api_whisper_model(custom_openai_base_url)
+        logging.debug("Will use whisper API on %s, %s",
+                      custom_openai_base_url, self.whisper_api_model)
 
     def transcribe(self) -> List[Segment]:
         logging.debug(
@@ -103,7 +107,7 @@ def transcribe(self) -> List[Segment]:
     def get_segments_for_file(self, file: str, offset_ms: int = 0):
         with open(file, "rb") as file:
             options = {
-                "model": "whisper-1",
+                "model": self.whisper_api_model,
                 "file": file,
                 "response_format": "verbose_json",
                 "prompt": self.transcription_task.transcription_options.initial_prompt,

diff --git a/buzz/transcriber/recording_transcriber.py b/buzz/transcriber/recording_transcriber.py
@@ -15,7 +15,7 @@
 from PyQt6.QtCore import QObject, pyqtSignal
 
 from buzz import transformers_whisper, whisper_audio
-from buzz.model_loader import ModelType
+from buzz.model_loader import ModelType, get_custom_api_whisper_model
 from buzz.settings.settings import Settings
 from buzz.transcriber.transcriber import TranscriptionOptions, Task
 from buzz.transcriber.whisper_cpp import WhisperCpp, whisper_cpp_params
@@ -55,6 +55,7 @@ def __init__(
         self.mutex = threading.Lock()
         self.sounddevice = sounddevice
         self.openai_client = None
+        self.whisper_api_model = get_custom_api_whisper_model("")
 
     def start(self):
         model_path = self.model_path
@@ -72,10 +73,13 @@ def start(self):
             custom_openai_base_url = settings.value(
                 key=Settings.Key.CUSTOM_OPENAI_BASE_URL, default_value=""
             )
+            self.whisper_api_model = get_custom_api_whisper_model(custom_openai_base_url)
             self.openai_client = OpenAI(
                 api_key=self.transcription_options.openai_access_token,
                 base_url=custom_openai_base_url if custom_openai_base_url else None
             )
+            logging.debug("Will use whisper API on %s, %s",
+                          custom_openai_base_url, self.whisper_api_model)
         else:  # ModelType.HUGGING_FACE
             model = transformers_whisper.load_model(model_path)
 
@@ -180,7 +184,7 @@ def start(self):
 
                             with open(temp_filename, 'rb') as temp_file:
                                 options = {
-                                    "model": "whisper-1",
+                                    "model": self.whisper_api_model,
                                     "file": temp_file,
                                     "response_format": "verbose_json",
                                     "prompt": self.transcription_options.initial_prompt,

diff --git a/buzz/widgets/recording_transcriber_widget.py b/buzz/widgets/recording_transcriber_widget.py
@@ -93,9 +93,7 @@ def __init__(
         if selected_model is None or selected_model.model_type not in model_types:
             selected_model = default_model
 
-        openai_access_token = ""
-        if selected_model.model_type == ModelType.OPEN_AI_WHISPER_API:
-            openai_access_token = get_password(key=Key.OPENAI_API_KEY)
+        openai_access_token = get_password(key=Key.OPENAI_API_KEY)
 
         self.transcription_options = TranscriptionOptions(
             model=selected_model,