Update app.py

app.py

@@ -3,8 +3,7 @@ from fastapi.responses import HTMLResponse
 from pydantic import BaseModel, Field
 import torch
 import re
-from nltk.tokenize import sent_tokenize
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from fastapi.middleware.cors import CORSMiddleware
 import asyncio
 from concurrent.futures import ThreadPoolExecutor
@@ -14,26 +13,17 @@ import time
 from typing import List, Optional
 import os
 import uvicorn
-import nltk
-try:
-    nltk.data.find('tokenizers/punkt_tab')
-except LookupError:
-    nltk.download('punkt_tab')
+import spacy
+from spacy.lang.en import English
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Download required NLTK data
-try:
-    nltk.data.find('tokenizers/punkt')
-except LookupError:
-    nltk.download('punkt', quiet=True)
-
 app = FastAPI(
     title="English to Telugu Translation API",
     description="High-performance translation service powered by MBart",
-    version="1.0.0"
+    version="2.0.0"
 )
 
 # Add CORS middleware
@@ -49,12 +39,14 @@ app.add_middleware(
 translator = None
 device = None
 executor = None
+nlp = None
+translation_pipeline = None
 
 # Pydantic models
 class TranslationRequest(BaseModel):
     text: str = Field(..., max_length=5000, min_length=1)
-    batch_size: Optional[int] = Field(default=8, ge=1, le=32)
-    max_length: Optional[int] = Field(default=512, ge=1, le=1024)
+    batch_size: Optional[int] = Field(default=16, ge=1, le=64)
+    max_length: Optional[int] = Field(default=256, ge=1, le=512)
 
 class TranslationResponse(BaseModel):
     original_text: str
@@ -62,110 +54,157 @@ class TranslationResponse(BaseModel):
     processing_time: float
     model_used: str
 
-@lru_cache(maxsize=2000)
-def cached_sentence_tokenize(text: str) -> tuple:
-    """Cached sentence tokenization"""
-    return tuple(sent_tokenize(text))
+# Fast sentence splitter using spaCy
+def get_fast_sentence_splitter():
+    """Initialize fast sentence splitter"""
+    try:
+        nlp = English()
+        nlp.add_pipe('sentencizer')
+        return nlp
+    except:
+        # Fallback to simple regex splitting
+        return None
+
+@lru_cache(maxsize=5000)
+def cached_sentence_split(text: str) -> tuple:
+    """Ultra-fast cached sentence splitting"""
+    if nlp:
+        doc = nlp(text)
+        return tuple([sent.text.strip() for sent in doc.sents if sent.text.strip()])
+    else:
+        # Fast regex fallback
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+        return tuple([s.strip() for s in sentences if s.strip()])
 
-class OptimizedTranslator:
+class UltraFastTranslator:
     def __init__(self, model, tokenizer, device):
         self.model = model
         self.tokenizer = tokenizer
         self.device = device
         
-        if hasattr(self.model, 'to'):
-            self.model = self.model.to(device)
-            logger.info(f"Model moved to device: {device}")
+        # Move model to device and set to eval mode
+        self.model = self.model.to(device).eval()
+        
+        # Enable optimizations
+        if hasattr(torch, 'jit') and device.type == 'cuda':
+            try:
+                self.model = torch.jit.script(self.model)
+                logger.info("Model optimized with TorchScript")
+            except:
+                logger.info("TorchScript optimization failed, using regular model")
+        
+        # Set model to half precision if using GPU
+        if device.type == 'cuda':
+            try:
+                self.model = self.model.half()
+                logger.info("Model converted to half precision")
+            except:
+                logger.info("Half precision conversion failed")
     
-    def translate_batch(self, sentences: List[str], max_length: int = 512) -> List[str]:
-        """Batch translation for better performance"""
+    def translate_batch_ultra_fast(self, sentences: List[str], max_length: int = 256) -> List[str]:
+        """Ultra-fast batch translation with aggressive optimizations"""
         if not sentences:
             return []
         
-        processed_sentences = []
-        for sentence in sentences:
+        # Filter and prepare sentences
+        valid_sentences = []
+        indices = []
+        
+        for i, sentence in enumerate(sentences):
             sentence = sentence.strip()
-            if sentence and not sentence.endswith(('.', '?', '!', ':', ';')):
-                sentence += '.'
-            processed_sentences.append(sentence)
+            if sentence:
+                if not sentence.endswith(('.', '?', '!', ':', ';')):
+                    sentence += '.'
+                valid_sentences.append(sentence)
+                indices.append(i)
+        
+        if not valid_sentences:
+            return [""] * len(sentences)
         
         try:
+            # Tokenize with aggressive truncation
             inputs = self.tokenizer(
-                processed_sentences,
+                valid_sentences,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
-                max_length=max_length
+                max_length=min(max_length, 128)  # Aggressive truncation
             )
             
             inputs = {k: v.to(self.device) for k, v in inputs.items()}
             
             with torch.no_grad():
+                # Ultra-fast generation settings
                 outputs = self.model.generate(
                     **inputs,
-                    max_length=max_length,
-                    num_beams=2,
+                    max_length=min(max_length, 128),
+                    num_beams=1,  # Greedy decoding (fastest)
                     early_stopping=True,
                     do_sample=False,
                     pad_token_id=self.tokenizer.pad_token_id,
-                    eos_token_id=self.tokenizer.eos_token_id
+                    eos_token_id=self.tokenizer.eos_token_id,
+                    use_cache=True
                 )
             
+            # Fast decoding
             translations = []
             for output in outputs:
-                translated_text = self.tokenizer.decode(output, skip_special_tokens=True)
+                translated_text = self.tokenizer.decode(
+                    output, 
+                    skip_special_tokens=True,
+                    clean_up_tokenization_spaces=False
+                )
                 translations.append(translated_text)
             
-            return translations
+            # Map back to original indices
+            result = [""] * len(sentences)
+            for i, translation in enumerate(translations):
+                result[indices[i]] = translation
             
-        except Exception as e:
-            logger.error(f"Batch translation error: {e}")
-            return [self.translate_single(sentence, max_length) for sentence in processed_sentences]
-    
-    def translate_single(self, sentence: str, max_length: int = 512) -> str:
-        """Single sentence translation fallback"""
-        try:
-            inputs = self.tokenizer(sentence, return_tensors="pt", max_length=max_length, truncation=True)
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            return result
             
-            with torch.no_grad():
-                outputs = self.model.generate(
-                    **inputs,
-                    max_length=max_length,
-                    num_beams=2,
-                    early_stopping=True,
-                    do_sample=False
-                )
-            
-            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         except Exception as e:
-            logger.error(f"Single translation error: {e}")
-            return sentence
+            logger.error(f"Batch translation error: {e}")
+            return [sentence for sentence in sentences]  # Return original on error
 
 @app.on_event("startup")
 async def load_models():
-    global translator, device, executor
+    global translator, device, executor, nlp, translation_pipeline
     
     start_time = time.time()
     logger.info("🚀 Loading translation models...")
     
+    # Use GPU if available
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     logger.info(f"📱 Using device: {device}")
     
-    executor = ThreadPoolExecutor(max_workers=2)
+    # Increase thread pool for better concurrency
+    executor = ThreadPoolExecutor(max_workers=4)
+    
+    # Initialize fast sentence splitter
+    nlp = get_fast_sentence_splitter()
     
     try:
         model_name = "aryaumesh/english-to-telugu"
         logger.info(f"📦 Loading model: {model_name}")
         
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        # Load with optimizations
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            use_fast=True  # Use fast tokenizer
+        )
+        
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16 if device.type == 'cuda' else torch.float32,
+            low_cpu_mem_usage=True
+        )
         
-        translator = OptimizedTranslator(model, tokenizer, device)
+        translator = UltraFastTranslator(model, tokenizer, device)
         
-        # Warm up the model
+        # Quick warm-up with small text
         logger.info("🔥 Warming up model...")
-        _ = translator.translate_single("Hello, this is a test.")
+        _ = translator.translate_batch_ultra_fast(["Hello."], max_length=64)
         
         load_time = time.time() - start_time
         logger.info(f"✅ Models loaded successfully in {load_time:.2f} seconds")
@@ -182,10 +221,11 @@ async def shutdown_event():
         executor.shutdown(wait=True)
     logger.info("✅ Shutdown complete")
 
-def process_translation_sync(text: str, batch_size: int = 8, max_length: int = 512) -> tuple:
-    """Synchronous translation processing"""
+def process_translation_ultra_fast(text: str, batch_size: int = 16, max_length: int = 256) -> tuple:
+    """Ultra-fast translation processing with aggressive optimizations"""
     start_time = time.time()
     
+    # Fast line processing
     lines = text.split('\n')
     translated_lines = []
     
@@ -195,16 +235,18 @@ def process_translation_sync(text: str, batch_size: int = 8, max_length: int = 5
             translated_lines.append("")
             continue
         
-        sentences = list(cached_sentence_tokenize(line))
+        # Use cached sentence splitting
+        sentences = list(cached_sentence_split(line))
         
         if not sentences:
             translated_lines.append("")
             continue
         
+        # Process in larger batches for better GPU utilization
         translated_sentences = []
         for i in range(0, len(sentences), batch_size):
             batch = sentences[i:i + batch_size]
-            batch_translations = translator.translate_batch(batch, max_length)
+            batch_translations = translator.translate_batch_ultra_fast(batch, max_length)
             translated_sentences.extend(batch_translations)
         
         translated_line = " ".join(translated_sentences)
@@ -217,7 +259,7 @@ def process_translation_sync(text: str, batch_size: int = 8, max_length: int = 5
 
 @app.post("/translate/", response_model=TranslationResponse)
 async def translate_text(request: TranslationRequest):
-    """Main translation endpoint"""
+    """Ultra-fast translation endpoint"""
     if not request.text.strip():
         return TranslationResponse(
             original_text=request.text,
@@ -230,17 +272,17 @@ async def translate_text(request: TranslationRequest):
         loop = asyncio.get_event_loop()
         translation, processing_time = await loop.run_in_executor(
             executor,
-            process_translation_sync,
+            process_translation_ultra_fast,
             request.text,
-            request.batch_size or 8,
-            request.max_length or 512
+            request.batch_size or 16,
+            min(request.max_length or 256, 256)  # Cap max length for speed
         )
         
         return TranslationResponse(
             original_text=request.text,
             translated_text=translation,
             processing_time=processing_time,
-            model_used="aryaumesh/english-to-telugu"
+            model_used="aryaumesh/english-to-telugu-optimized"
         )
         
     except Exception as e:
@@ -254,14 +296,14 @@ async def translate_text(request: TranslationRequest):
 
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
-    """Root endpoint with simple HTML interface"""
+    """Root endpoint with optimized HTML interface"""
     html_content = """
     <!DOCTYPE html>
     <html lang="en">
     <head>
         <meta charset="UTF-8">
         <meta name="viewport" content="width=device-width, initial-scale=1.0">
-        <title>English to Telugu Translation</title>
+        <title>Ultra-Fast English to Telugu Translation</title>
         <style>
             * { margin: 0; padding: 0; box-sizing: border-box; }
             body { 
@@ -279,13 +321,21 @@ async def read_root():
                 overflow: hidden;
             }
             .header {
-                background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%);
+                background: linear-gradient(135deg, #ff6b6b 0%, #feca57 100%);
                 color: white;
                 padding: 30px;
                 text-align: center;
             }
             .header h1 { font-size: 2.5em; margin-bottom: 10px; }
             .header p { font-size: 1.1em; opacity: 0.9; }
+            .speed-badge {
+                display: inline-block;
+                background: rgba(255,255,255,0.2);
+                padding: 5px 15px;
+                border-radius: 20px;
+                margin-top: 10px;
+                font-weight: bold;
+            }
             .content { padding: 40px; }
             .form-group { margin-bottom: 25px; }
             label { 
@@ -308,8 +358,8 @@ async def read_root():
             }
             textarea:focus {
                 outline: none;
-                border-color: #4facfe;
-                box-shadow: 0 0 0 3px rgba(79, 172, 254, 0.1);
+                border-color: #ff6b6b;
+                box-shadow: 0 0 0 3px rgba(255, 107, 107, 0.1);
             }
             .controls {
                 display: flex;
@@ -334,7 +384,7 @@ async def read_root():
                 width: 100px;
             }
             button { 
-                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                background: linear-gradient(135deg, #ff6b6b 0%, #feca57 100%);
                 color: white; 
                 padding: 15px 30px; 
                 border: none; 
@@ -347,7 +397,7 @@ async def read_root():
             }
             button:hover { 
                 transform: translateY(-2px);
-                box-shadow: 0 10px 20px rgba(102, 126, 234, 0.3);
+                box-shadow: 0 10px 20px rgba(255, 107, 107, 0.3);
             }
             button:disabled {
                 opacity: 0.7;
@@ -359,7 +409,7 @@ async def read_root():
                 padding: 25px; 
                 background: #f8f9ff; 
                 border-radius: 10px;
-                border-left: 4px solid #4facfe;
+                border-left: 4px solid #ff6b6b;
             }
             .result h3 { 
                 color: #333; 
@@ -389,7 +439,7 @@ async def read_root():
                 border: 1px solid #e0e0e0;
                 font-size: 0.9em;
             }
-            .stat strong { color: #4facfe; }
+            .stat strong { color: #ff6b6b; }
             .loading {
                 display: none;
                 text-align: center;
@@ -400,7 +450,7 @@ async def read_root():
                 width: 40px;
                 height: 40px;
                 border: 4px solid #f3f3f3;
-                border-top: 4px solid #4facfe;
+                border-top: 4px solid #ff6b6b;
                 border-radius: 50%;
                 animation: spin 1s linear infinite;
             }
@@ -408,6 +458,10 @@ async def read_root():
                 0% { transform: rotate(0deg); }
                 100% { transform: rotate(360deg); }
             }
+            .speed-indicator {
+                color: #27ae60;
+                font-weight: bold;
+            }
             .error {
                 background: #ffe6e6;
                 border-left-color: #ff4757;
@@ -426,31 +480,32 @@ async def read_root():
     <body>
         <div class="container">
             <div class="header">
-                <h1>🌐 English to Telugu Translation</h1>
-                <p>Powered by Advanced AI • Fast • Accurate</p>
+                <h1>⚡ Ultra-Fast English to Telugu Translation</h1>
+                <p>Lightning Speed AI Translation • Optimized for Performance</p>
+                <div class="speed-badge">🚀 Target: &lt;0.5s processing time</div>
             </div>
             
             <div class="content">
                 <div class="form-group">
                     <label for="inputText">📝 Enter English Text:</label>
-                    <textarea id="inputText" placeholder="Type or paste your English text here...&#10;&#10;You can enter multiple sentences or even paragraphs.&#10;The system will translate everything while preserving the structure."></textarea>
+                    <textarea id="inputText" placeholder="Type or paste your English text here...&#10;&#10;Optimized for ultra-fast processing!&#10;Try shorter sentences for best performance."></textarea>
                 </div>
                 
                 <div class="controls">
                     <div class="control-group">
                         <label>Batch Size:</label>
-                        <input type="number" id="batchSize" value="8" min="1" max="32">
+                        <input type="number" id="batchSize" value="16" min="1" max="64">
                     </div>
                     <div class="control-group">
                         <label>Max Length:</label>
-                        <input type="number" id="maxLength" value="512" min="1" max="1024">
+                        <input type="number" id="maxLength" value="256" min="1" max="512">
                     </div>
-                    <button onclick="translateText()">🚀 Translate</button>
+                    <button onclick="translateText()">⚡ Ultra-Fast Translate</button>
                 </div>
                 
                 <div class="loading" id="loading">
                     <div class="spinner"></div>
-                    <p>Translating your text...</p>
+                    <p>Processing at lightning speed...</p>
                 </div>
                 
                 <div id="result" class="result" style="display: none;">
@@ -458,7 +513,7 @@ async def read_root():
                     <div id="translatedText" class="translated-text"></div>
                     <div class="stats">
                         <div class="stat">
-                            <strong>Processing Time:</strong> <span id="processingTime">-</span> seconds
+                            <strong>Processing Time:</strong> <span id="processingTime" class="speed-indicator">-</span> seconds
                         </div>
                         <div class="stat">
                             <strong>Model:</strong> <span id="modelUsed">-</span>
@@ -474,8 +529,8 @@ async def read_root():
         <script>
             async function translateText() {
                 const inputText = document.getElementById('inputText').value;
-                const batchSize = parseInt(document.getElementById('batchSize').value) || 8;
-                const maxLength = parseInt(document.getElementById('maxLength').value) || 512;
+                const batchSize = parseInt(document.getElementById('batchSize').value) || 16;
+                const maxLength = parseInt(document.getElementById('maxLength').value) || 256;
                 
                 if (!inputText.trim()) {
                     alert('⚠️ Please enter some text to translate');
@@ -487,11 +542,13 @@ async def read_root():
                 const result = document.getElementById('result');
                 
                 // Show loading state
-                button.textContent = '⏳ Translating...';
+                button.textContent = '⚡ Processing...';
                 button.disabled = true;
                 loading.style.display = 'block';
                 result.style.display = 'none';
                 
+                const startTime = performance.now();
+                
                 try {
                     const response = await fetch('/translate/', {
                         method: 'POST',
@@ -510,10 +567,24 @@ async def read_root():
                     }
                     
                     const data = await response.json();
+                    const totalTime = (performance.now() - startTime) / 1000;
                     
                     // Display results
                     document.getElementById('translatedText').textContent = data.translated_text;
-                    document.getElementById('processingTime').textContent = data.processing_time.toFixed(2);
+                    
+                    const processingTimeElement = document.getElementById('processingTime');
+                    processingTimeElement.textContent = data.processing_time.toFixed(3);
+                    
+                    // Color code based on speed
+                    if (data.processing_time < 0.5) {
+                        processingTimeElement.style.color = '#27ae60'; // Green
+                        processingTimeElement.textContent += ' ⚡';
+                    } else if (data.processing_time < 2) {
+                        processingTimeElement.style.color = '#f39c12'; // Orange
+                    } else {
+                        processingTimeElement.style.color = '#e74c3c'; // Red
+                    }
+                    
                     document.getElementById('modelUsed').textContent = data.model_used;
                     document.getElementById('charCount').textContent = data.original_text.length;
                     
@@ -537,7 +608,7 @@ async def read_root():
                     result.classList.add('error');
                     
                 } finally {
-                    button.textContent = '🚀 Translate';
+                    button.textContent = '⚡ Ultra-Fast Translate';
                     button.disabled = false;
                 }
             }
@@ -568,6 +639,7 @@ async def health_check():
         "status": "healthy",
         "device": str(device) if device else "not_initialized",
         "model_loaded": translator is not None,
+        "optimization_level": "ultra_fast",
         "timestamp": time.time()
     }
 
@@ -575,21 +647,24 @@ async def health_check():
 async def api_info():
     """API information endpoint"""
     return {
-        "title": "English to Telugu Translation API",
-        "version": "1.0.0",
+        "title": "Ultra-Fast English to Telugu Translation API",
+        "version": "2.0.0",
         "model": "aryaumesh/english-to-telugu",
+        "optimizations": [
+            "spacy_sentence_splitting",
+            "aggressive_caching",
+            "gpu_half_precision",
+            "greedy_decoding",
+            "larger_batch_sizes",
+            "reduced_beam_search"
+        ],
+        "target_processing_time": "< 0.5 seconds",
         "endpoints": {
             "translate": "/translate/",
             "health": "/health",
             "docs": "/docs",
             "redoc": "/redoc"
-        },
-        "supported_features": [
-            "batch_processing",
-            "caching",
-            "async_processing",
-            "input_validation"
-        ]
+        }
     }
 
 if __name__ == "__main__":
@@ -600,4 +675,4 @@ if __name__ == "__main__":
         port=port,
         workers=1,
         log_level="info"
-    ) 
+    )