Browse Source

update models list

subDesTagesMitExtraKaese 1 month ago
parent
commit
3b255ce27e
3 changed files with 63 additions and 27 deletions
  1. 22 2
      .github/workflows/docker.yml
  2. 2 1
      Dockerfile
  3. 39 24
      speech_recognition.py

+ 22 - 2
.github/workflows/docker.yml

@@ -112,6 +112,26 @@ jobs:
           context: .
           platforms: linux/amd64,linux/arm64/v8
           push: true
-          tags: ftcaplan/matrix-stt-bot:large-q5_0
+          tags: ftcaplan/matrix-stt-bot:large-v3-q5_0
           build-args: |
-            "PRELOAD_MODEL=large-q5_0"
+            "PRELOAD_MODEL=large-v3-q5_0"
+      -
+        name: Build and push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          platforms: linux/amd64,linux/arm64/v8
+          push: true
+          tags: ftcaplan/matrix-stt-bot:large-v3-turbo
+          build-args: |
+            "PRELOAD_MODEL=large-v3-turbo"
+      -
+        name: Build and push
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          platforms: linux/amd64,linux/arm64/v8
+          push: true
+          tags: ftcaplan/matrix-stt-bot:large-v3-turbo-q5_0
+          build-args: |
+            "PRELOAD_MODEL=large-v3-turbo-q5_0"

+ 2 - 1
Dockerfile

@@ -31,9 +31,10 @@ COPY --from=builder /build/main /app/
 VOLUME /data/
 
 ADD ./*.py /app/
+ADD ./whisper.cpp/models/download-ggml-model.sh /app/
 
 ARG PRELOAD_MODEL
 ENV PRELOAD_MODEL ${PRELOAD_MODEL}
-RUN if [ -n "$PRELOAD_MODEL" ]; then wget -nv "https://ggml.ggerganov.com/ggml-model-whisper-$PRELOAD_MODEL.bin"; fi
+RUN if [ -n "$PRELOAD_MODEL" ]; then /app/download-ggml-model.sh "$PRELOAD_MODEL" "/app"; fi
 
 CMD ["python3", "-u", "main.py"]

+ 39 - 24
speech_recognition.py

@@ -28,21 +28,36 @@ def convert_audio(data: bytes, out_filename: str):
   return out
 
 MODELS = [
-  "tiny.en", 
-  "tiny.en-q5_1", 
-  "tiny", 
-  "tiny-q5_1", 
-  "base.en", 
-  "base.en-q5_1", 
-  "base", 
-  "base-q5_1", 
-  "small.en", 
-  "small.en-q5_1", 
-  "small", 
+  "tiny",
+  "tiny.en",
+  "tiny-q5_1",
+  "tiny.en-q5_1",
+  "tiny-q8_0",
+  "base",
+  "base.en",
+  "base-q5_1",
+  "base.en-q5_1",
+  "base-q8_0",
+  "small",
+  "small.en",
+  "small.en-tdrz",
   "small-q5_1",
-  "medium.en-q5_0", 
-  "medium-q5_0", 
-  "large-q5_0"
+  "small.en-q5_1",
+  "small-q8_0",
+  "medium",
+  "medium.en",
+  "medium-q5_0",
+  "medium.en-q5_0",
+  "medium-q8_0",
+  "large-v1",
+  "large-v2",
+  "large-v2-q5_0",
+  "large-v2-q8_0",
+  "large-v3",
+  "large-v3-q5_0",
+  "large-v3-turbo",
+  "large-v3-turbo-q5_0",
+  "large-v3-turbo-q8_0",
 ]
 
 class ASR():
@@ -52,20 +67,20 @@ class ASR():
     self.model = model
     self.language = language
 
-    if os.path.exists(f"/app/ggml-model-whisper-{model}.bin"):
-      self.model_path = f"/app/ggml-model-whisper-{model}.bin"
+    if os.path.exists(f"/app/ggml-{model}.bin"):
+      self.model_path = f"/app"
     else:
-      self.model_path = f"/data/models/ggml-{model}.bin"
-      if not os.path.exists("/data/models"):
-        os.mkdir("/data/models")
-        
-    self.model_url = f"https://ggml.ggerganov.com/ggml-model-whisper-{self.model}.bin"
+      self.model_path = f"/data/models"
+      if not os.path.exists(self.model_path):
+        os.mkdir(self.model_path)
+
     self.lock = asyncio.Lock()
 
   def load_model(self):
-    if not os.path.exists(self.model_path) or os.path.getsize(self.model_path) == 0:
+    file_path = f"{self.model_path}/ggml-{self.model}.bin"
+    if not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
       print("Downloading model...")
-      subprocess.run(["wget", "-nv", self.model_url, "-O", self.model_path], check=True)
+      subprocess.run(["./download-ggml-model.sh", self.model, self.model_path], check=True)
       print("Done.")
 
   async def transcribe(self, audio: bytes) -> str:
@@ -74,7 +89,7 @@ class ASR():
     async with self.lock:
       proc = await asyncio.create_subprocess_exec(
           "./main",
-          "-m", self.model_path,
+          "-m", f"{self.model_path}/ggml-{self.model}.bin",
           "-l", self.language,
           "-f", filename,
           "-nt",