Update app.py

app.py

@@ -1,17 +1,28 @@
 from fastapi import FastAPI, Request
-from transformers import pipeline
 from fastapi.responses import JSONResponse
+from transformers import pipeline
 import uvicorn
 
 app = FastAPI()
-generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct", max_new_tokens=300)
+
+# Load your model
+generator = pipeline(
+    "text-generation",
+    model="mistralai/Mistral-7B-Instruct",
+    max_new_tokens=400,
+    temperature=0.7,
+    do_sample=True
+)
 
 @app.post("/ask")
-async def ask_ai(request: Request):
+async def ask_voxtreme(request: Request):
     data = await request.json()
     prompt = data.get("prompt", "")
+
     if not prompt:
-        return {"error": "No prompt given"}
+        return JSONResponse(content={"error": "Prompt is missing"}, status_code=400)
+
+    # Generate text using model
+    output = generator(prompt)[0]["generated_text"]
+    return JSONResponse(content={"response": output})
 
-    result = generator(prompt, do_sample=True, temperature=0.7)
-    return JSONResponse(content={"response": result[0]['generated_text']})