Update app.py

app.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel
 import torch
 from transformers import (
@@ -13,21 +14,26 @@ import uuid
 import os
 from typing import Optional
 
-# Initialize FastAPI app
-app = FastAPI()
+# FastAPI instance
+app = FastAPI(title="Kinyarwanda NLP API", version="1.0")
 
-# Configuration
+# Config
 MODEL_PATH = "/app/models/suno-bark"
 SENTIMENT_MODEL_PATH = "/app/models/sentiment"
+SAMPLE_RATE = 24000
 
-# Load all models in a single try-except block
+# Ensure working directory for audio
+AUDIO_DIR = "/tmp/audio"
+os.makedirs(AUDIO_DIR, exist_ok=True)
+
+# Load models
 try:
-    # TTS Model
+    # TTS
     tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
     processor = AutoProcessor.from_pretrained(MODEL_PATH)
     model = BarkModel.from_pretrained(MODEL_PATH)
-    
-    # Sentiment Analysis Model (pre-downloaded)
+
+    # Sentiment
     sentiment_tokenizer = AutoTokenizer.from_pretrained(SENTIMENT_MODEL_PATH)
     sentiment_model = AutoModelForSequenceClassification.from_pretrained(SENTIMENT_MODEL_PATH)
     sentiment_pipeline = pipeline(
@@ -37,15 +43,15 @@ try:
         truncation=True,
         max_length=512
     )
-    
-    # Device configuration
+
+    # Device config
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
-    
+
 except Exception as e:
-    raise RuntimeError(f"Initialization failed: {str(e)}")
+    raise RuntimeError(f"Model initialization failed: {e}")
 
-# Request models
+# Request schemas
 class TTSRequest(BaseModel):
     text: str
 
@@ -56,63 +62,55 @@ class LegalDocRequest(BaseModel):
     text: str
     domain: Optional[str] = "general"
 
+# Root route
 @app.get("/")
 def root():
     return {"message": "Welcome to Kinyarwanda NLP API"}
 
+# Text-to-Speech endpoint
 @app.post("/tts/")
 def text_to_speech(request: TTSRequest):
-    output_file = None
+    output_file = os.path.join(AUDIO_DIR, f"tts_{uuid.uuid4().hex}.wav")
+
     try:
         inputs = processor(request.text, return_tensors="pt").to(device)
         with torch.no_grad():
-            speech = model.generate(**inputs)
-        
-        output_file = f"output_{uuid.uuid4().hex}.wav"
-        wavfile.write(output_file, rate=24000, data=speech.cpu().numpy().squeeze())
-        
+            audio_array = model.generate(**inputs)
+
+        wavfile.write(output_file, rate=SAMPLE_RATE, data=audio_array.cpu().numpy().squeeze())
+
         return FileResponse(
             output_file,
             media_type="audio/wav",
-            filename=output_file
+            filename=os.path.basename(output_file)
         )
-    
+
     except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={"error": f"TTS failed: {str(e)}"}
-        )
-    
+        raise HTTPException(status_code=500, detail=f"TTS generation failed: {str(e)}")
+
     finally:
-        if output_file and os.path.exists(output_file):
+        if os.path.exists(output_file):
             os.remove(output_file)
 
+# Sentiment Analysis endpoint
 @app.post("/sentiment/")
 def analyze_sentiment(request: SentimentRequest):
     try:
         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)}"}
-        )
+        raise HTTPException(status_code=500, detail=f"Sentiment analysis failed: {str(e)}")
 
+# Legal Parsing endpoint
 @app.post("/legal-parse/")
 def parse_legal_document(request: LegalDocRequest):
     try:
         keywords = ["contract", "agreement", "party", "terms", "confidential", "jurisdiction"]
-        found_keywords = [kw for kw in keywords if kw in request.text.lower()]
-        
+        found = [kw for kw in keywords if kw in request.text.lower()]
         return {
-            "identified_keywords": found_keywords,
+            "identified_keywords": found,
             "domain": request.domain,
             "status": "success"
         }
-        
     except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={"error": f"Legal parsing failed: {str(e)}"}
-        )
+        raise HTTPException(status_code=500, detail=f"Legal parsing failed: {str(e)}")