Update app.py

app.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, File, UploadFile, Form
 from fastapi import HTTPException
 import uvicorn
 from pydantic import BaseModel
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 import torch
 import ffmpeg
 import io
@@ -19,6 +19,14 @@ model_name = "openai/whisper-base"
 device = 0 if torch.cuda.is_available() else -1  # Use GPU if available
 whisper_pipeline = pipeline("automatic-speech-recognition", model=model_name, device=device)
 
+# set up translation pipeline
+def get_translation_pipeline(translation_model_path):
+    model = AutoModelForSeq2SeqLM.from_pretrained(translation_model_path)
+    tokenizer = AutoTokenizer.from_pretrained(translation_model_path)
+    translation_pipeline = pipeline('translation', model=model, tokenizer=tokenizer, device=device)
+    return translation_pipeline
+
+translator = get_translation_pipeline("mutisya/nllb_600m-en-kik-kam-luo-mer-som-swh-drL-24_5-filtered-v24_28_4")
 
 class RecognitionResponse(BaseModel):
     text: str
@@ -76,6 +84,27 @@ async def recognize_audio(audio: UploadFile = File(...), language: str = Form("e
             logger.error("Unexpected error during transcription", exc_info=True)
             raise HTTPException(status_code=500, detail="Internal Server Error")
 
+class TranslationResponse(BaseModel):
+    translated_text: str
+
+@app.post("/translate", response_model=TranslationResponse)
+async def translate_text(request: TranslationRequest):
+    source_language = request.source_language
+    target_language = request.target_language
+    text_to_translate = request.text
+    
+    try:
+        src_lang = flores_codes[source_language]
+        tgt_lang = flores_codes[target_language]
+
+        translated_text = translator(text_to_translate, src_lang=src_lang,tgt_lang=tgt_lang)[0]['translation_text']
+        
+        return TranslationResponse(translated_text=translated_text)
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+        
+
 # Run the FastAPI application
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)