Update app.py

app.py

@@ -1,35 +1,39 @@
-from transformers import AutoTokenizer, AutoProcessor
 from fastapi import FastAPI, Form
 from fastapi.responses import FileResponse, JSONResponse
 from pydantic import BaseModel
 import torch
-from transformers import AutoProcessor, BarkModel, pipeline
+from transformers import AutoTokenizer, AutoProcessor, BarkModel, pipeline
 import scipy.io.wavfile as wavfile
 import uuid
 import os
 from typing import Optional
 
-# Load the pre-downloaded model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("/app/models/suno-bark")
-processor = AutoProcessor.from_pretrained("/app/models/suno-bark")
-
-# Rest of your application code...
-
-# Load TTS model and processor
-processor = AutoProcessor.from_pretrained("suno/bark")
-model = BarkModel.from_pretrained("suno/bark")
-
-# Load sentiment analysis pipeline (using multilingual model)
-sentiment_model = pipeline("sentiment-analysis", model="nlptown/bert-base-multilingual-uncased-sentiment")
-
-# Ensure model is on CPU or CUDA if available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Ensure proper model loading from pre-downloaded path
+MODEL_PATH = "/app/models/suno-bark"
+
+# Load models and processors once during startup
+try:
+    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="nlptown/bert-base-multilingual-uncased-sentiment"
+    )
+    
+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)
 
-# FastAPI app
+# Initialize FastAPI app
 app = FastAPI()
 
-# Endpoint input models
+# Request models
 class TTSRequest(BaseModel):
     text: str
 
@@ -47,20 +51,20 @@ def root():
 @app.post("/tts/")
 def text_to_speech(request: TTSRequest):
     try:
-        # Generate speech
-        inputs = processor(request.text, return_tensors="pt")
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        speech = model.generate(**inputs)
-
-        # Save audio
+        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"
-        speech_np = speech.cpu().numpy().squeeze()
-        wavfile.write(output_file, rate=22050, data=speech_np)
-
-        return FileResponse(output_file, media_type="audio/wav")
-
+        wavfile.write(output_file, rate=24000, data=speech.cpu().numpy().squeeze())
+        
+        return FileResponse(output_file, media_type="audio/wav", filename=output_file)
+    
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": f"TTS failed: {str(e)}"})
+    finally:
+        if os.path.exists(output_file):
+            os.remove(output_file)
 
 @app.post("/sentiment/")
 def analyze_sentiment(request: SentimentRequest):
@@ -68,14 +72,19 @@ def analyze_sentiment(request: SentimentRequest):
         result = sentiment_model(request.text)
         return {"result": result}
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": 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:
-        # Placeholder logic (replace with training-based custom logic)
-        keywords = ["contract", "agreement", "party", "terms"]
+        # 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()]
-        return {"identified_keywords": found_keywords, "domain": request.domain}
+        
+        return {
+            "identified_keywords": found_keywords,
+            "domain": request.domain,
+            "status": "success"
+        }
     except Exception as e:
-        return JSONResponse(status_code=500, content={"error": str(e)})
+        return JSONResponse(status_code=500, content={"error": f"Legal parsing failed: {str(e)}"})