Update app.py

app.py

@@ -121,20 +121,29 @@ def remove_repetitive_responses(responses):
     return unique_responses
 
 @app.post("/generate/")
+@GPU(duration=0)
 async def generate(request: ChatRequest):
     try:
-        normalized_message = normalize_input(request.message)
-        with ThreadPoolExecutor() as executor:
-            futures = [executor.submit(model.generate, f"<s>[INST]{normalized_message} [/INST]",
-                                       top_k=request.top_k, top_p=request.top_p, temperature=request.temperature) 
-                       for model in global_data['models'].values()]
-            responses = [{'model': model, 'response': future.result()} 
-                         for model, future in zip(global_data['models'].keys(), as_completed(futures))]
-        
+        inputs = normalize_input(request.message)
+        futures = [
+            executor.submit(model.generate, inputs, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
+            for model in global_data['models'].values()
+        ]
+        responses = [{'model': model, 'response': future.result()} for model, future in zip(global_data['models'].keys(), as_completed(futures))]
         unique_responses = remove_repetitive_responses(responses)
         return unique_responses
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating responses: {e}")
+        print(f"Error generating responses: {e}")
+        raise HTTPException(status_code=500, detail="Error generating responses")
+
+@app.middleware("http")
+async def process_request(request: Request, call_next):
+    try:
+        response = await call_next(request)
+        return response
+    except Exception as e:
+        print(f"Request error: {e}")
+        raise HTTPException(status_code=500, detail="Internal Server Error")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)