Create app.py

app.py

@@ -0,0 +1,128 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from IndicTransToolkit import IndicProcessor
+
+# Initialize FastAPI app
+app = FastAPI(
+    title="Indic Translation API",
+    description="API for translating text between English and Indic languages",
+    version="1.0.0"
+)
+
+# Define request body model
+class InputData(BaseModel):
+    sentences: List[str]
+    target_lang: str
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "sentences": ["Hello, how are you?", "What is your name?"],
+                "target_lang": "hin_Deva"
+            }
+        }
+
+# Initialize models and processors
+try:
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        "ai4bharat/indictrans2-en-indic-1B",
+        trust_remote_code=True
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        "ai4bharat/indictrans2-en-indic-1B",
+        trust_remote_code=True
+    )
+    ip = IndicProcessor(inference=True)
+    DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(DEVICE)
+except Exception as e:
+    raise RuntimeError(f"Failed to load models: {str(e)}")
+
+@app.get("/")
+async def root():
+    """Root endpoint returning API information"""
+    return {
+        "message": "Welcome to the Indic Translation API",
+        "status": "active",
+        "supported_languages": [
+            "hin_Deva",  # Hindi
+            "ben_Beng",  # Bengali
+            "tam_Taml",  # Tamil
+            # Add other supported languages here
+        ]
+    }
+
+@app.post("/translate/")
+async def translate(input_data: InputData):
+    """
+    Translate text from English to specified Indic language
+    
+    Args:
+        input_data: InputData object containing sentences and target language
+    
+    Returns:
+        Dictionary containing translated text
+    """
+    try:
+        # Source language is always English
+        src_lang = "eng_Latn"
+        tgt_lang = input_data.target_lang
+
+        # Preprocess the input sentences
+        batch = ip.preprocess_batch(
+            input_data.sentences,
+            src_lang=src_lang,
+            tgt_lang=tgt_lang
+        )
+
+        # Tokenize the sentences
+        inputs = tokenizer(
+            batch,
+            truncation=True,
+            padding="longest",
+            return_tensors="pt",
+            return_attention_mask=True
+        ).to(DEVICE)
+
+        # Generate translations
+        with torch.no_grad():
+            generated_tokens = model.generate(
+                **inputs,
+                use_cache=True,
+                min_length=0,
+                max_length=256,
+                num_beams=5,
+                num_return_sequences=1
+            )
+
+        # Decode the generated tokens
+        with tokenizer.as_target_tokenizer():
+            generated_tokens = tokenizer.batch_decode(
+                generated_tokens.detach().cpu().tolist(),
+                skip_special_tokens=True,
+                clean_up_tokenization_spaces=True
+            )
+
+        # Postprocess the translations
+        translations = ip.postprocess_batch(generated_tokens, lang=tgt_lang)
+
+        return {
+            "translations": translations,
+            "source_language": src_lang,
+            "target_language": tgt_lang
+        }
+
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=f"Translation error: {str(e)}"
+        )
+
+# Add health check endpoint
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {"status": "healthy"}