fix-endpoint

src/server/main.py

@@ -21,14 +21,27 @@ from IndicTransToolkit import IndicProcessor
 from logging_config import logger
 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
+from typing import Annotated, Any, OrderedDict, List
 import zipfile
 import soundfile as sf
+import torch
+from fastapi import Body, FastAPI, HTTPException, Response
+from parler_tts import ParlerTTSForConditionalGeneration
+from transformers import AutoTokenizer, AutoFeatureExtractor, set_seed
 import numpy as np
 from config import SPEED, ResponseFormat, config
+from logger import logger
+import uvicorn
+import argparse
+from fastapi.responses import RedirectResponse, StreamingResponse
+import io
+import os
+import logging
 
 # Device setup
 if torch.cuda.is_available():
@@ -76,13 +89,29 @@ class TTSModelManager:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         description_tokenizer = AutoTokenizer.from_pretrained(model.config.text_encoder._name_or_path)
 
+        # Set pad tokens
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         if description_tokenizer.pad_token is None:
             description_tokenizer.pad_token = description_tokenizer.eos_token
+
         
-        model.forward = torch.compile(model.forward, mode="default")
+        # TODO - temporary disable -torch.compile 
+        '''
+        # Update model configuration
+        model.config.pad_token_id = tokenizer.pad_token_id
+        # Update for deprecation: use max_batch_size instead of batch_size
+        if hasattr(model.generation_config.cache_config, 'max_batch_size'):
+            model.generation_config.cache_config.max_batch_size = 1
+        model.generation_config.cache_implementation = "static"
+        '''
+        # Compile the model
+        compile_mode = "default"
+        #compile_mode = "reduce-overhead"
+        
+        model.forward = torch.compile(model.forward, mode=compile_mode)
 
+        # Warmup
         warmup_inputs = tokenizer("Warmup text for compilation", 
                                 return_tensors="pt", 
                                 padding="max_length", 
@@ -95,7 +124,8 @@ class TTSModelManager:
             "prompt_attention_mask": warmup_inputs["attention_mask"],
         }
         
-        for _ in range(1):
+        n_steps = 1 if compile_mode == "default" else 2
+        for _ in range(n_steps):
             _ = model.generate(**model_kwargs)
         
         logger.info(
@@ -122,14 +152,16 @@ 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 multiple languages",
+    description="AI Chat API supporting Indian languages",
     version="1.0.0",
     redirect_slashes=False,
     lifespan=lifespan
 )
 
+
 def chunk_text(text, chunk_size):
     words = text.split()
     chunks = []
@@ -165,6 +197,7 @@ 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"]
@@ -290,23 +323,14 @@ async def generate_audio_batch(
 
     return StreamingResponse(in_memory_zip, media_type="application/zip")
 
+
 # Supported language codes
 SUPPORTED_LANGUAGES = {
-    # Indian languages
     "asm_Beng", "kas_Arab", "pan_Guru", "ben_Beng", "kas_Deva", "san_Deva",
     "brx_Deva", "mai_Deva", "sat_Olck", "doi_Deva", "mal_Mlym", "snd_Arab",
     "eng_Latn", "mar_Deva", "snd_Deva", "gom_Deva", "mni_Beng", "tam_Taml",
     "guj_Gujr", "mni_Mtei", "tel_Telu", "hin_Deva", "npi_Deva", "urd_Arab",
-    "kan_Knda", "ory_Orya",
-    # European languages
-    "deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn",
-    "por_Latn", "rus_Cyrl", "pol_Latn"
-}
-
-# Define European languages for direct processing
-EUROPEAN_LANGUAGES = {
-    "deu_Latn", "fra_Latn", "nld_Latn", "spa_Latn", "ita_Latn",
-    "por_Latn", "rus_Cyrl", "pol_Latn"
+    "kan_Knda", "ory_Orya"
 }
 
 class Settings(BaseSettings):
@@ -328,6 +352,7 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
+
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -341,6 +366,7 @@ app.state.limiter = limiter
 
 llm_manager = LLMManager(settings.llm_model_name)
 
+# Translation Manager and Model Manager
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 class TranslateManager:
@@ -356,7 +382,7 @@ class TranslateManager:
         elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
             model_name = "ai4bharat/indictrans2-indic-indic-dist-320M" if use_distilled else "ai4bharat/indictrans2-indic-indic-1B"
         else:
-            raise ValueError("Invalid language combination: English to English or European languages not supported here.")
+            raise ValueError("Invalid language combination: English to English translation is not supported.")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -389,7 +415,7 @@ class ModelManager:
         elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
             key = 'indic_indic'
         else:
-            raise ValueError("Invalid language combination for translation.")
+            raise ValueError("Invalid language combination: English to English translation is not supported.")
 
         if key not in self.models:
             if self.is_lazy_loading:
@@ -406,10 +432,11 @@ class ModelManager:
 ip = IndicProcessor(inference=True)
 model_manager = ModelManager()
 
+# Pydantic Models
 class ChatRequest(BaseModel):
     prompt: str
-    src_lang: str = "kan_Knda"
-    tgt_lang: str = "kan_Knda"
+    src_lang: str = "kan_Knda"  # Default to Kannada
+    tgt_lang: str = "kan_Knda"  # Default to Kannada
 
     @field_validator("prompt")
     def prompt_must_be_valid(cls, v):
@@ -434,9 +461,11 @@ class TranslationRequest(BaseModel):
 class TranslationResponse(BaseModel):
     translations: List[str]
 
+# Dependency to get TranslateManager
 def get_translate_manager(src_lang: str, tgt_lang: str) -> TranslateManager:
     return model_manager.get_model(src_lang, tgt_lang)
 
+# Internal Translation Endpoint
 @app.post("/translate", response_model=TranslationResponse)
 async def translate(request: TranslationRequest, translate_manager: TranslateManager = Depends(get_translate_manager)):
     input_sentences = request.sentences
@@ -476,12 +505,14 @@ async def translate(request: TranslationRequest, translate_manager: TranslateMan
     translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
     return TranslationResponse(translations=translations)
 
+# Helper function to perform internal translation
 async def perform_internal_translation(sentences: List[str], src_lang: str, tgt_lang: str) -> List[str]:
     translate_manager = model_manager.get_model(src_lang, tgt_lang)
     request = TranslationRequest(sentences=sentences, src_lang=src_lang, tgt_lang=tgt_lang)
     response = await translate(request, translate_manager)
     return response.translations
 
+# API Endpoints
 @app.get("/v1/health")
 async def health_check():
     return {"status": "healthy", "model": settings.llm_model_name}
@@ -533,14 +564,14 @@ async def chat(request: Request, chat_request: ChatRequest):
     if not chat_request.prompt:
         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:
-        # Determine if the language requires translation (Indian languages only)
-        is_indian_language = chat_request.src_lang not in EUROPEAN_LANGUAGES and chat_request.src_lang != "eng_Latn"
-        is_target_indian = chat_request.tgt_lang not in EUROPEAN_LANGUAGES and chat_request.tgt_lang != "eng_Latn"
-
-        if is_indian_language:
-            # Translate prompt to English for Indian languages
+        # 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,
@@ -553,12 +584,13 @@ async def chat(request: Request, chat_request: ChatRequest):
             prompt_to_process = chat_request.prompt
             logger.info("Prompt in English or European language, no translation needed")
 
-        # Generate response directly with the LLM
+        # 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}")
 
-        if is_target_indian and chat_request.tgt_lang != "eng_Latn":
-            # Translate response to target Indian language
+        # 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",
@@ -588,10 +620,8 @@ async def visual_query(
         if image.size == (0, 0):
             raise HTTPException(status_code=400, detail="Uploaded image is empty or invalid")
 
-        is_indian_language = src_lang not in EUROPEAN_LANGUAGES and src_lang != "eng_Latn"
-        is_target_indian = tgt_lang not in EUROPEAN_LANGUAGES and tgt_lang != "eng_Latn"
-
-        if is_indian_language:
+        # Translate query to English if src_lang is not English
+        if src_lang != "eng_Latn":
             translated_query = await perform_internal_translation(
                 sentences=[query],
                 src_lang=src_lang,
@@ -601,12 +631,14 @@ async def visual_query(
             logger.info(f"Translated query to English: {query_to_process}")
         else:
             query_to_process = query
-            logger.info("Query in English or European language, no translation needed")
+            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}")
 
-        if is_target_indian and tgt_lang != "eng_Latn":
+        # Translate answer to target language if tgt_lang is not English
+        if tgt_lang != "eng_Latn":
             translated_answer = await perform_internal_translation(
                 sentences=[answer],
                 src_lang="eng_Latn",
@@ -616,7 +648,7 @@ async def visual_query(
             logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
         else:
             final_answer = answer
-            logger.info(f"Answer in {tgt_lang}, no translation needed")
+            logger.info("Answer kept in English, no translation needed")
 
         return {"answer": final_answer}
     except Exception as e:
@@ -640,16 +672,14 @@ async def chat_v2(
     logger.info(f"Received prompt: {prompt}, src_lang: {src_lang}, tgt_lang: {tgt_lang}, Image provided: {image is not None}")
 
     try:
-        is_indian_language = src_lang not in EUROPEAN_LANGUAGES and src_lang != "eng_Latn"
-        is_target_indian = tgt_lang not in EUROPEAN_LANGUAGES and tgt_lang != "eng_Latn"
-
         if image:
             image_data = await image.read()
             if not image_data:
                 raise HTTPException(status_code=400, detail="Uploaded image is empty")
             img = Image.open(io.BytesIO(image_data))
 
-            if is_indian_language:
+            # Translate prompt to English if src_lang is not English
+            if src_lang != "eng_Latn":
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
@@ -659,12 +689,13 @@ async def chat_v2(
                 logger.info(f"Translated prompt to English: {prompt_to_process}")
             else:
                 prompt_to_process = prompt
-                logger.info("Prompt in English or European language, no translation needed")
+                logger.info("Prompt already in English, no translation needed")
 
             decoded = await llm_manager.chat_v2(img, prompt_to_process)
-            logger.info(f"Generated response: {decoded}")
+            logger.info(f"Generated English response: {decoded}")
 
-            if is_target_indian and tgt_lang != "eng_Latn":
+            # Translate response to target language if tgt_lang is not English
+            if tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
@@ -674,9 +705,10 @@ async def chat_v2(
                 logger.info(f"Translated response to {tgt_lang}: {final_response}")
             else:
                 final_response = decoded
-                logger.info(f"Response in {tgt_lang}, no translation needed")
+                logger.info("Response kept in English, no translation needed")
         else:
-            if is_indian_language:
+            # Translate prompt to English if src_lang is not English
+            if src_lang != "eng_Latn":
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
@@ -686,12 +718,13 @@ async def chat_v2(
                 logger.info(f"Translated prompt to English: {prompt_to_process}")
             else:
                 prompt_to_process = prompt
-                logger.info("Prompt in English or European language, no translation needed")
+                logger.info("Prompt already in English, no translation needed")
 
             decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-            logger.info(f"Generated response: {decoded}")
+            logger.info(f"Generated English response: {decoded}")
 
-            if is_target_indian and tgt_lang != "eng_Latn":
+            # Translate response to target language if tgt_lang is not English
+            if tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
@@ -701,7 +734,7 @@ async def chat_v2(
                 logger.info(f"Translated response to {tgt_lang}: {final_response}")
             else:
                 final_response = decoded
-                logger.info(f"Response in {tgt_lang}, no translation needed")
+                logger.info("Response kept in English, no translation needed")
 
         return ChatResponse(response=final_response)
     except Exception as e:
@@ -711,6 +744,7 @@ async def chat_v2(
 class TranscriptionResponse(BaseModel):
     text: str
 
+
 class ASRModelManager:
     def __init__(self, device_type="cuda"):
         self.device_type = device_type
@@ -722,25 +756,54 @@ class ASRModelManager:
             "telugu": "te", "urdu": "ur"
         }
 
+
+from fastapi import FastAPI, UploadFile
+import torch
+import torchaudio
+from transformers import AutoModel
+import argparse
+import uvicorn
+from pydantic import BaseModel
+from pydub import AudioSegment
+from fastapi import FastAPI, File, UploadFile, HTTPException, Query
+from fastapi.responses import RedirectResponse, JSONResponse
+from typing import List
+
+# Load the model
 model = AutoModel.from_pretrained("ai4bharat/indic-conformer-600m-multilingual", trust_remote_code=True)
-asr_manager = ASRModelManager()
+
+asr_manager = ASRModelManager()  # Load Kannada, Hindi, Tamil, Telugu, Malayalam
+
+
+#asr_manager = ASRModelManager(device_type="")
 
 @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)
 
-    target_sample_rate = 16000
+    # 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 CTC decoding
+    #transcription_ctc = model(wav, "kn", "ctc")
+
+    # Perform ASR with RNNT decoding
     transcription_rnnt = model(wav, "kn", "rnnt")
+
     return JSONResponse(content={"text": transcription_rnnt})
 
+
+
 class BatchTranscriptionResponse(BaseModel):
     transcriptions: List[str]
 
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run the FastAPI server.")
     parser.add_argument("--port", type=int, default=settings.port, help="Port to run the server on.")