Update app.py

app.py

@@ -1,5 +1,4 @@
-from fastapi import FastAPI, Form
-from fastapi.responses import FileResponse, JSONResponse
+from fastapi import FastAPI
 from pydantic import BaseModel
 import torch
 from transformers import AutoTokenizer, AutoProcessor, BarkModel, pipeline
@@ -8,30 +7,34 @@ import uuid
 import os
 from typing import Optional
 
-# Ensure proper model loading from pre-downloaded path
+# Initialize FastAPI app
+app = FastAPI()
+
+# Configuration
 MODEL_PATH = "/app/models/suno-bark"
+SENTIMENT_MODEL = "cardiffnlp/twitter-xlm-roberta-base-sentiment"  # PyTorch-compatible model
 
-# Load models and processors once during startup
+# Load all models in a single try-except block
 try:
+    # TTS Model
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     processor = AutoProcessor.from_pretrained(MODEL_PATH)
     model = BarkModel.from_pretrained(MODEL_PATH)
     
-    # Load sentiment analysis pipeline
-    sentiment_model = pipeline(
+    # Sentiment Analysis Model
+    sentiment_pipeline = pipeline(
         "sentiment-analysis",
-        model="nlptown/bert-base-multilingual-uncased-sentiment"
+        model=SENTIMENT_MODEL,
+        truncation=True,
+        max_length=512
     )
     
+    # Device configuration
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    
 except Exception as e:
-    raise RuntimeError(f"Model loading failed: {str(e)}")
-
-# Device configuration
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Initialize FastAPI app
-app = FastAPI()
+    raise RuntimeError(f"Initialization failed: {str(e)}")
 
 # Request models
 class TTSRequest(BaseModel):
@@ -46,10 +49,11 @@ class LegalDocRequest(BaseModel):
 
 @app.get("/")
 def root():
-    return {"message": "Welcome to Kinyarwanda NLP API"}
+    return {"message": "Welcome to Kinyarwanda-Engine"}
 
 @app.post("/tts/")
 def text_to_speech(request: TTSRequest):
+    output_file = None
     try:
         inputs = processor(request.text, return_tensors="pt").to(device)
         with torch.no_grad():
@@ -58,26 +62,37 @@ def text_to_speech(request: TTSRequest):
         output_file = f"output_{uuid.uuid4().hex}.wav"
         wavfile.write(output_file, rate=24000, data=speech.cpu().numpy().squeeze())
         
-        return FileResponse(output_file, media_type="audio/wav", filename=output_file)
+        return FileResponse(
+            output_file,
+            media_type="audio/wav",
+            filename=output_file
+        )
     
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"TTS failed: {str(e)}"})
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"TTS failed: {str(e)}"}
+        )
+    
     finally:
-        if os.path.exists(output_file):
+        if output_file and os.path.exists(output_file):
             os.remove(output_file)
 
 @app.post("/sentiment/")
 def analyze_sentiment(request: SentimentRequest):
     try:
-        result = sentiment_model(request.text)
+        result = sentiment_pipeline(request.text)
         return {"result": result}
+        
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"Sentiment analysis failed: {str(e)}"})
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Sentiment analysis failed: {str(e)}"}
+        )
 
 @app.post("/legal-parse/")
 def parse_legal_document(request: LegalDocRequest):
     try:
-        # Basic keyword extraction (replace with trained model in production)
         keywords = ["contract", "agreement", "party", "terms", "confidential", "jurisdiction"]
         found_keywords = [kw for kw in keywords if kw in request.text.lower()]
         
@@ -86,5 +101,9 @@ def parse_legal_document(request: LegalDocRequest):
             "domain": request.domain,
             "status": "success"
         }
+        
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": f"Legal parsing failed: {str(e)}"})
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Legal parsing failed: {str(e)}"}
+        )