add-uer

src/server/main.py

@@ -21,27 +21,14 @@ 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, List
+from typing import Annotated, Any, OrderedDict
 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():
@@ -89,29 +76,13 @@ 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
-
         
-        # 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)
+        model.forward = torch.compile(model.forward, mode="default")
 
-        # Warmup
         warmup_inputs = tokenizer("Warmup text for compilation", 
                                 return_tensors="pt", 
                                 padding="max_length", 
@@ -124,8 +95,7 @@ class TTSModelManager:
             "prompt_attention_mask": warmup_inputs["attention_mask"],
         }
         
-        n_steps = 1 if compile_mode == "default" else 2
-        for _ in range(n_steps):
+        for _ in range(1):
             _ = model.generate(**model_kwargs)
         
         logger.info(
@@ -152,16 +122,14 @@ 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",
+    description="AI Chat API supporting multiple languages",
     version="1.0.0",
     redirect_slashes=False,
     lifespan=lifespan
 )
 
-
 def chunk_text(text, chunk_size):
     words = text.split()
     chunks = []
@@ -197,7 +165,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,14 +290,23 @@ 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"
+    "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"
 }
 
 class Settings(BaseSettings):
@@ -352,7 +328,6 @@ class Settings(BaseSettings):
 
 settings = Settings()
 
-
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -366,7 +341,6 @@ 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:
@@ -382,7 +356,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 translation is not supported.")
+            raise ValueError("Invalid language combination: English to English or European languages not supported here.")
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
         model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -415,7 +389,7 @@ class ModelManager:
         elif not src_lang.startswith("eng") and not tgt_lang.startswith("eng"):
             key = 'indic_indic'
         else:
-            raise ValueError("Invalid language combination: English to English translation is not supported.")
+            raise ValueError("Invalid language combination for translation.")
 
         if key not in self.models:
             if self.is_lazy_loading:
@@ -432,11 +406,10 @@ class ModelManager:
 ip = IndicProcessor(inference=True)
 model_manager = ModelManager()
 
-# Pydantic Models
 class ChatRequest(BaseModel):
     prompt: str
-    src_lang: str = "kan_Knda"  # Default to Kannada
-    tgt_lang: str = "kan_Knda"  # Default to Kannada
+    src_lang: str = "kan_Knda"
+    tgt_lang: str = "kan_Knda"
 
     @field_validator("prompt")
     def prompt_must_be_valid(cls, v):
@@ -461,11 +434,9 @@ 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
@@ -505,14 +476,12 @@ 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}
@@ -564,9 +533,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}")
+
     try:
-        # Translate prompt to English if src_lang is not English
-        if chat_request.src_lang != "eng_Latn":
+        # 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
             translated_prompt = await perform_internal_translation(
                 sentences=[chat_request.prompt],
                 src_lang=chat_request.src_lang,
@@ -575,15 +549,16 @@ 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 already in English, no translation needed")
+            logger.info("Prompt in English or European language, no translation needed")
 
-        # Generate response in English
+        # Generate response directly with the LLM
         response = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-        logger.info(f"Generated English response: {response}")
+        logger.info(f"Generated response: {response}")
 
-        # Translate response to target language if tgt_lang is not English
-        if chat_request.tgt_lang != "eng_Latn":
+        if is_target_indian and chat_request.tgt_lang != "eng_Latn":
+            # Translate response to target Indian language
             translated_response = await perform_internal_translation(
                 sentences=[response],
                 src_lang="eng_Latn",
@@ -592,8 +567,9 @@ 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("Response kept in English, no translation needed")
+            logger.info(f"Response in {chat_request.tgt_lang}, no translation needed")
 
         return ChatResponse(response=final_response)
     except Exception as e:
@@ -612,8 +588,10 @@ 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":
+        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:
             translated_query = await perform_internal_translation(
                 sentences=[query],
                 src_lang=src_lang,
@@ -623,14 +601,12 @@ async def visual_query(
             logger.info(f"Translated query to English: {query_to_process}")
         else:
             query_to_process = query
-            logger.info("Query already in English, no translation needed")
+            logger.info("Query in English or European language, 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":
+        if is_target_indian and tgt_lang != "eng_Latn":
             translated_answer = await perform_internal_translation(
                 sentences=[answer],
                 src_lang="eng_Latn",
@@ -640,7 +616,7 @@ async def visual_query(
             logger.info(f"Translated answer to {tgt_lang}: {final_answer}")
         else:
             final_answer = answer
-            logger.info("Answer kept in English, no translation needed")
+            logger.info(f"Answer in {tgt_lang}, no translation needed")
 
         return {"answer": final_answer}
     except Exception as e:
@@ -664,14 +640,16 @@ 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))
 
-            # Translate prompt to English if src_lang is not English
-            if src_lang != "eng_Latn":
+            if is_indian_language:
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
@@ -681,13 +659,12 @@ async def chat_v2(
                 logger.info(f"Translated prompt to English: {prompt_to_process}")
             else:
                 prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
+                logger.info("Prompt in English or European language, no translation needed")
 
             decoded = await llm_manager.chat_v2(img, prompt_to_process)
-            logger.info(f"Generated English response: {decoded}")
+            logger.info(f"Generated response: {decoded}")
 
-            # Translate response to target language if tgt_lang is not English
-            if tgt_lang != "eng_Latn":
+            if is_target_indian and tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
@@ -697,10 +674,9 @@ async def chat_v2(
                 logger.info(f"Translated response to {tgt_lang}: {final_response}")
             else:
                 final_response = decoded
-                logger.info("Response kept in English, no translation needed")
+                logger.info(f"Response in {tgt_lang}, no translation needed")
         else:
-            # Translate prompt to English if src_lang is not English
-            if src_lang != "eng_Latn":
+            if is_indian_language:
                 translated_prompt = await perform_internal_translation(
                     sentences=[prompt],
                     src_lang=src_lang,
@@ -710,13 +686,12 @@ async def chat_v2(
                 logger.info(f"Translated prompt to English: {prompt_to_process}")
             else:
                 prompt_to_process = prompt
-                logger.info("Prompt already in English, no translation needed")
+                logger.info("Prompt in English or European language, no translation needed")
 
             decoded = await llm_manager.generate(prompt_to_process, settings.max_tokens)
-            logger.info(f"Generated English response: {decoded}")
+            logger.info(f"Generated response: {decoded}")
 
-            # Translate response to target language if tgt_lang is not English
-            if tgt_lang != "eng_Latn":
+            if is_target_indian and tgt_lang != "eng_Latn":
                 translated_response = await perform_internal_translation(
                     sentences=[decoded],
                     src_lang="eng_Latn",
@@ -726,7 +701,7 @@ async def chat_v2(
                 logger.info(f"Translated response to {tgt_lang}: {final_response}")
             else:
                 final_response = decoded
-                logger.info("Response kept in English, no translation needed")
+                logger.info(f"Response in {tgt_lang}, no translation needed")
 
         return ChatResponse(response=final_response)
     except Exception as e:
@@ -736,7 +711,6 @@ async def chat_v2(
 class TranscriptionResponse(BaseModel):
     text: str
 
-
 class ASRModelManager:
     def __init__(self, device_type="cuda"):
         self.device_type = device_type
@@ -748,54 +722,25 @@ 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()  # Load Kannada, Hindi, Tamil, Telugu, Malayalam
-
-
-#asr_manager = ASRModelManager(device_type="")
+asr_manager = ASRModelManager()
 
 @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
+    target_sample_rate = 16000
     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.")