Update main.py

main.py

@@ -1,28 +1,25 @@
 from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
-from enum import Enum
-from transformers import pipeline, MarianMTModel, MarianTokenizer
-import shutil
-import os
-import uuid
-from googletrans import Translator
+from pydantic import BaseModel
+from transformers import pipeline, MarianMTModel, MarianTokenizer, WhisperProcessor, WhisperForConditionalGeneration
+import torch
+import tempfile
+import soundfile as sf
 
 app = FastAPI()
 
-# 🌍 Language Enum for dropdown in Swagger
-class LanguageEnum(str, Enum):
-    ta = "ta"  # Tamil
-    fr = "fr"  # French
-    es = "es"  # Spanish
-    de = "de"  # German
-    it = "it"  # Italian
-    hi = "hi"  # Hindi
-    ru = "ru"  # Russian
-    zh = "zh"  # Chinese
-    ar = "ar"  # Arabic
-
-# 🌐 Map target language to translation model
-model_map = {
+# Allow frontend to call backend
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Supported languages
+translation_models = {
     "fr": "Helsinki-NLP/opus-mt-en-fr",
     "es": "Helsinki-NLP/opus-mt-en-es",
     "de": "Helsinki-NLP/opus-mt-en-de",
@@ -31,78 +28,108 @@ model_map = {
     "ru": "Helsinki-NLP/opus-mt-en-ru",
     "zh": "Helsinki-NLP/opus-mt-en-zh",
     "ar": "Helsinki-NLP/opus-mt-en-ar",
-    "ta": "gsarti/opus-mt-en-ta"
+    "ta": "Helsinki-NLP/opus-mt-en-ta"
 }
 
-def translate_text(text, target_lang):
-    if target_lang == "ta":
-        try:
-            translator = Translator()
-            result = translator.translate(text, dest="ta")
-            return result.text
-        except Exception as e:
-            return f"Google Translate failed: {str(e)}"
-
-    if target_lang not in model_map:
-        return f"No model for language: {target_lang}"
-
-    model_name = model_map[target_lang]
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    encoded = tokenizer([text], return_tensors="pt", padding=True)
-    translated = model.generate(**encoded)
-    return tokenizer.batch_decode(translated, skip_special_tokens=True)[0]
-
-# 🧠 Generate a random English sentence
-def generate_random_sentence(prompt="Daily conversation", max_length=30):
-    generator = pipeline("text-generation", model="distilgpt2")
-    result = generator(prompt, max_length=max_length, num_return_sequences=1)
-    return result[0]["generated_text"].strip()
-
-# 🎤 Transcription endpoint
-@app.post("/transcribe")
-async def transcribe(audio: UploadFile = File(...)):
-    temp_filename = f"temp_{uuid.uuid4().hex}.wav"
-    with open(temp_filename, "wb") as f:
-        shutil.copyfileobj(audio.file, f)
+# Load models once
+generator = pipeline("text-generation", model="gpt2")
+whisper_model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-base")
+whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-base")
+
+@app.get("/")
+def root():
+    return {"message": "Backend is live ✅"}
+
+@app.get("/generate")
+def generate_and_translate(prompt: str, target_lang: str):
     try:
-        asr = pipeline("automatic-speech-recognition", model="openai/whisper-medium")
-        result = asr(temp_filename)
-        return JSONResponse(content={"transcribed_text": result["text"]})
-    finally:
-        os.remove(temp_filename)
+        if target_lang not in translation_models:
+            return {"error": "Unsupported language."}
+        
+        # 1. Generate English sentence
+        result = generator(prompt, max_length=30, num_return_sequences=1)[0]["generated_text"]
+        english_sentence = result.strip()
+
+        # 2. Translate
+        model_name = translation_models[target_lang]
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        tokens = tokenizer(english_sentence, return_tensors="pt", padding=True)
+        translated_ids = model.generate(**tokens)
+        translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+
+        return {"english": english_sentence, "translated": translated_text}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
+
+class TranslateRequest(BaseModel):
+    text: str
+    target_lang: str
 
-# 🌍 Translation endpoint
 @app.post("/translate")
-async def translate(text: str = Form(...), target_lang: LanguageEnum = Form(...)):
-    translated = translate_text(text, target_lang.value)
-    return JSONResponse(content={"translated_text": translated})
+def translate_text(data: TranslateRequest):
+    try:
+        if data.target_lang not in translation_models:
+            return {"error": "Unsupported language."}
+        
+        model_name = translation_models[data.target_lang]
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        tokens = tokenizer(data.text, return_tensors="pt", padding=True)
+        translated_ids = model.generate(**tokens)
+        translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+
+        return {"translated_text": translated_text}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
+
+@app.post("/transcribe")
+async def transcribe_audio(audio: UploadFile = File(...)):
+    try:
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        temp_file.write(await audio.read())
+        temp_file.close()
+
+        audio_data, _ = sf.read(temp_file.name)
+        inputs = whisper_processor(audio_data, sampling_rate=16000, return_tensors="pt")
+        predicted_ids = whisper_model.generate(inputs["input_features"])
+        transcription = whisper_processor.decode(predicted_ids[0], skip_special_tokens=True)
+
+        return {"transcribed_text": transcription}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})
 
-# 🔁 Combined endpoint (speech-to-translation)
 @app.post("/process")
-async def process(audio: UploadFile = File(...), target_lang: LanguageEnum = Form(...)):
-    temp_filename = f"temp_{uuid.uuid4().hex}.wav"
-    with open(temp_filename, "wb") as f:
-        shutil.copyfileobj(audio.file, f)
+async def transcribe_and_translate_audio(
+    audio: UploadFile = File(...),
+    target_lang: str = Form(...)
+):
     try:
-        asr = pipeline("automatic-speech-recognition", model="openai/whisper-medium")
-        result = asr(temp_filename)
-        transcribed_text = result["text"]
-        translated_text = translate_text(transcribed_text, target_lang.value)
-        return JSONResponse(content={
-            "transcribed_text": transcribed_text,
-            "translated_text": translated_text
-        })
-    finally:
-        os.remove(temp_filename)
+        if target_lang not in translation_models:
+            return {"error": "Unsupported language."}
 
-# ✨ Generate + Translate endpoint
-@app.get("/generate")
-def generate(prompt: str = "Daily conversation", target_lang: LanguageEnum = LanguageEnum.it):
-    english = generate_random_sentence(prompt)
-    translated = translate_text(english, target_lang.value)
-    return {
-        "prompt": prompt,
-        "english": english,
-        "translated": translated
-    }
+        # Save uploaded file
+        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav")
+        temp_file.write(await audio.read())
+        temp_file.close()
+
+        # Transcribe
+        audio_data, _ = sf.read(temp_file.name)
+        inputs = whisper_processor(audio_data, sampling_rate=16000, return_tensors="pt")
+        predicted_ids = whisper_model.generate(inputs["input_features"])
+        transcription = whisper_processor.decode(predicted_ids[0], skip_special_tokens=True)
+
+        # Translate
+        model_name = translation_models[target_lang]
+        tokenizer = MarianTokenizer.from_pretrained(model_name)
+        model = MarianMTModel.from_pretrained(model_name)
+        tokens = tokenizer(transcription, return_tensors="pt", padding=True)
+        translated_ids = model.generate(**tokens)
+        translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
+
+        return {
+            "transcribed_text": transcription,
+            "translated_text": translated_text
+        }
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": str(e)})