update: minors corrections

.dockerignore

@@ -0,0 +1 @@
+venv

Dockerfile

@@ -1,5 +1,5 @@
 FROM python:3.10-slim
-# FROM pytorch/pytorch:2.1.0-cuda11.8-cudnn8-runtime
+# FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04
 
 WORKDIR /app
 
@@ -7,8 +7,14 @@ RUN apt-get update && apt-get install -y git ffmpeg
 
 COPY . .
 
+# RUN pip install --upgrade pip setuptools wheel
+
 RUN pip install --no-cache-dir -r requirements.txt
 
 # EXPOSE 8000
 # CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+
+# docker buildx build --platform linux/amd64 --no-cache -t ai-api .
+# docker tag ai-api kora3/ai-api:latest
+# docker push kora3/ai-api:latest

functions/translation.py

@@ -1,23 +1,25 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
-MODEL_NAME = "facebook/nllb-200-distilled-600M"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+MODEL_NAME = "facebook/nllb-200-distilled-600M" # facebook/nllb-200-3.3B
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
 model.eval()
 
 
 def translateText(text: str, sourceLang: str, targetLang: str) -> str:
     tokenizer.src_lang = sourceLang
-    inputs = tokenizer(text, return_tensors="pt")
+    inputs = tokenizer(text, return_tensors="pt").to(device)
 
     with torch.no_grad():
         tokens = model.generate(
             **inputs,
             forced_bos_token_id=tokenizer.convert_tokens_to_ids(targetLang),
-            num_beams=1,
-            max_length=128
+            max_new_tokens=512,
+            num_beams=4,
+            early_stopping=True,
         )
 
     translatedText = tokenizer.batch_decode(tokens, skip_special_tokens=True)[0]

language/dioula/dyu_stt.py

@@ -1,9 +1,11 @@
-import base64, tempfile, os
+import base64, tempfile, os, torch
 from transformers import pipeline
 from functions.utils import getAudioDuration
 
 MODEL_NAME = "facebook/mms-1b-all"
-pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "dyu"})
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "dyu"}, device=device)
 
 
 def dioulaSTT(audioBase64: str) -> dict:
@@ -20,4 +22,4 @@ def dioulaSTT(audioBase64: str) -> dict:
     finally:
         os.remove(tempAudioPath)
 
-    return {'text': text, 'language': 'mos', 'duration': duration}
+    return {'text': text, 'language': 'dyu', 'duration': duration}

language/moore/mos_stt.py

@@ -1,10 +1,12 @@
-import base64, tempfile, os
+import base64, tempfile, os, torch
 from transformers import pipeline
 from functions.utils import getAudioDuration
 # from huggingface_hub import login
 
 MODEL_NAME = "facebook/mms-1b-all"
-pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "mos"})
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, model_kwargs={"target_lang": "mos"}, device=device)
 # login(token=os.environ["HF_TOKEN"])
 # MODEL_NAME = "burkimbia/BIA-WHISPER-LARGE-SACHI_V3"
 # pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME)

main.py

@@ -2,8 +2,10 @@ from fastapi import FastAPI, Request, HTTPException
 from functions.translation import translateText
 from functions.speech_to_text import speechToText
 from functions.text_to_speech import textToSpeech
+
 from language.moore.mos_stt import mooreSTT
 from language.moore.mos_tts import mooreTTS
+
 from language.dioula.dyu_stt import dioulaSTT
 from language.dioula.dyu_tts import dioulaTTS