fix-bud

src/server/main.py

@@ -23,7 +23,6 @@ from tts_config import SPEED, ResponseFormat, config as tts_config
 from gemma_llm import LLMManager
 # from auth import get_api_key, settings as auth_settings
 
-
 import time
 from contextlib import asynccontextmanager
 from typing import Annotated, Any, OrderedDict, List
@@ -95,7 +94,6 @@ class TTSModelManager:
         if description_tokenizer.pad_token is None:
             description_tokenizer.pad_token = description_tokenizer.eos_token
 
-        
         # TODO - temporary disable -torch.compile 
         '''
         # Update model configuration
@@ -152,7 +150,6 @@ async def lifespan(_: FastAPI):
         tts_model_manager.get_or_load_model(config.model)
     yield
 
-#app = FastAPI(lifespan=lifespan)
 app = FastAPI(
     title="Dhwani API",
     description="AI Chat API supporting Indian languages",
@@ -161,7 +158,6 @@ app = FastAPI(
     lifespan=lifespan
 )
 
-
 def chunk_text(text, chunk_size):
     words = text.split()
     chunks = []
@@ -197,7 +193,6 @@ async def generate_audio(
                                 padding="max_length",
                                 max_length=tts_model_manager.max_length).to(device)
         
-        # Use the tensor fields directly instead of BatchEncoding object
         input_ids = desc_inputs["input_ids"]
         attention_mask = desc_inputs["attention_mask"]
         prompt_input_ids = prompt_inputs["input_ids"]
@@ -323,7 +318,6 @@ async def generate_audio_batch(
 
     return StreamingResponse(in_memory_zip, media_type="application/zip")
 
-
 # Supported language codes
 SUPPORTED_LANGUAGES = {
     "asm_Beng", "kas_Arab", "pan_Guru", "ben_Beng", "kas_Deva", "san_Deva",
@@ -354,7 +348,6 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -543,7 +536,7 @@ async def load_all_models():
         return {"status": "success", "message": "All models loaded"}
     except Exception as e:
         logger.error(f"Error loading models: {str(e)}")
-        raise HTTPException(status_code=500, detail=f"Failed to load models: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Failed to unload models: {str(e)}")
 
 @app.post("/v1/translate", response_model=TranslationResponse)
 async def translate_endpoint(request: TranslationRequest):
@@ -567,13 +560,10 @@ async def chat(request: Request, chat_request: ChatRequest):
         raise HTTPException(status_code=400, detail="Prompt cannot be empty")
     logger.info(f"Received prompt: {chat_request.prompt}, src_lang: {chat_request.src_lang}, tgt_lang: {chat_request.tgt_lang}")
     
-    # Define European languages that gemma-3-4b-it can handle natively
     EUROPEAN_LANGUAGES = {"deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn", "por_Latn", "rus_Cyrl", "pol_Latn"}
     
     try:
-        # Check if the source language is Indian (requires translation) or European/English (direct processing)
         if chat_request.src_lang != "eng_Latn" and chat_request.src_lang not in EUROPEAN_LANGUAGES:
-            # Translate Indian language prompt to English
             translated_prompt = await perform_internal_translation(
                 sentences=[chat_request.prompt],
                 src_lang=chat_request.src_lang,
@@ -582,17 +572,13 @@ async def chat(request: Request, chat_request: ChatRequest):
             prompt_to_process = translated_prompt[0]
             logger.info(f"Translated prompt to English: {prompt_to_process}")
         else:
-            # Use prompt directly for English and European languages
             prompt_to_process = chat_request.prompt
             logger.info("Prompt in English or European language, no translation needed")
 
-        # Generate response with the LLM (assumed to handle multilingual input natively)
         response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
         logger.info(f"Generated response: {response}")
 
-        # Check if the target language is Indian (requires translation) or European/English (direct output)
         if chat_request.tgt_lang != "eng_Latn" and chat_request.tgt_lang not in EUROPEAN_LANGUAGES:
-            # Translate response to Indian target language
             translated_response = await perform_internal_translation(
                 sentences=[response],
                 src_lang="eng_Latn",
@@ -601,7 +587,6 @@ async def chat(request: Request, chat_request: ChatRequest):
             final_response = translated_response[0]
             logger.info(f"Translated response to {chat_request.tgt_lang}: {final_response}")
         else:
-            # Keep response as-is for English and European languages
             final_response = response
             logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
 
@@ -622,7 +607,6 @@ async def visual_query(
         if image.size == (0, 0):
             raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
 
-        # Translate query to English if src_lang is not English
         if src_lang != "eng_Latn":
             translated_query = await perform_internal_translation(
                 sentences=[query],
@@ -635,11 +619,9 @@ async def visual_query(
             query_to_process = query
             logger.info("Query already in English, no translation needed")
 
-        # Generate response in English
         answer = await llm_manager.vision_query(image, query_to_process)
         logger.info(f"Generated English answer: {answer}")
 
-        # Translate answer to target language if tgt_lang is not English
         if tgt_lang != "eng_Latn":
             translated_answer = await perform_internal_translation(
                 sentences=[answer],
@@ -680,7 +662,6 @@ async def chat_v2(
                 raise HTTPException(status_code=400, detail="Uploaded image is empty")
             img = Image.open(io.BytesIO(image_data))
 
-            # Translate prompt to English if src_lang is not English
             if src_lang != "eng_Latn":
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
@@ -696,7 +677,6 @@ async def chat_v2(
             decoded = await llm_manager.chat_v2(img, prompt_to_process)
             logger.info(f"Generated English response: {decoded}")
 
-            # Translate response to target language if tgt_lang is not English
             if tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
@@ -709,7 +689,6 @@ async def chat_v2(
                 final_response = decoded
                 logger.info("Response kept in English, no translation needed")
         else:
-            # Translate prompt to English if src_lang is not English
             if src_lang != "eng_Latn":
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
@@ -725,7 +704,6 @@ async def chat_v2(
             decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
             logger.info(f"Generated English response: {decoded}")
 
-            # Translate response to target language if tgt_lang is not English
             if tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
@@ -746,7 +724,6 @@ async def chat_v2(
 class TranscriptionResponse(BaseModel):
     text: str
 
-
 class ASRModelManager:
     def __init__(self, device_type="cuda"):
         self.device_type = device_type
@@ -758,7 +735,6 @@ class ASRModelManager:
             "telugu": "te", "urdu": "ur"
         }
 
-
 from fastapi import FastAPI, UploadFile
 import torch
 import torchaudio
@@ -774,27 +750,30 @@ from typing import List
 # Load the model
 model = AutoModel.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True)
 
-asr_manager = ASRModelManager()  # Load Kannada, Hindi, Tamil, Telugu, Malayalam
+asr_manager = ASRModelManager()
 
-
-#asr_manager = ASRModelManager(device_type="")
+# Language to script mapping
+LANGUAGE_TO_SCRIPT = {
+    "kannada": "kan_Knda", "hindi": "hin_Deva", "malayalam": "mal_Mlym", "tamil": "tam_Taml",
+    "telugu": "tel_Telu", "assamese": "asm_Beng", "bengali": "ben_Beng", "gujarati": "guj_Gujr",
+    "marathi": "mar_Deva", "odia": "ory_Orya", "punjabi": "pan_Guru", "urdu": "urd_Arab",
+    # Add more as needed
+}
 
 @app.post("/transcribe/", response_model=TranscriptionResponse)
 async def transcribe_audio(file: UploadFile = File(...), language: str = Query(..., enum=list(asr_manager.model_language.keys()))):
-    # Load the uploaded audio file
-    wav, sr = torchaudio.load(file.file)
-    wav = torch.mean(wav, dim=0, keepdim=True)
-
-    # Resample if necessary
-    target_sample_rate = 16000  # Expected sample rate
-    if sr != target_sample_rate:
-        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate)
-        wav = resampler(wav)
-
-    # Perform ASR with RNNT decoding using the provided language
-    transcription_rnnt = model(wav, asr_manager.model_language[language], "rnnt")
-
-    return JSONResponse(content={"text": transcription_rnnt})
+    try:
+        wav, sr = torchaudio.load(file.file)
+        wav = torch.mean(wav, dim=0, keepdim=True)
+        target_sample_rate = 16000
+        if sr != target_sample_rate:
+            resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=target_sample_rate)
+            wav = resampler(wav)
+        transcription_rnnt = model(wav, asr_manager.model_language[language], "rnnt")
+        return TranscriptionResponse(text=transcription_rnnt)
+    except Exception as e:
+        logger.error(f"Error in transcription: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
 
 @app.post("/v1/speech_to_speech")
 async def speech_to_speech(
@@ -809,8 +788,8 @@ async def speech_to_speech(
     # Step 2: Process text with chat endpoint
     chat_request = ChatRequest(
         prompt=transcription.text,
-        src_lang="kn_Knda",  # Assuming script for Indian languages
-        tgt_lang="kn_Knda"
+        src_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda"),  # Dynamic script mapping
+        tgt_lang=LANGUAGE_TO_SCRIPT.get(language, "kan_Knda")
     )
     processed_text = await chat(Request(), chat_request)
     logger.info(f"Processed text: {processed_text.response}")
@@ -825,7 +804,6 @@ async def speech_to_speech(
     )
     return audio_response
 
-
 class BatchTranscriptionResponse(BaseModel):
     transcriptions: List[str]