Update app.py

app.py

@@ -1,4 +1,3 @@
-
 from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from llama_cpp import Llama
@@ -6,12 +5,19 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import uvicorn
 import re
 from dotenv import load_dotenv
-import spaces
+from spaces import GPU
 
 load_dotenv()
 
 app = FastAPI()
 
+# Initialize ZeroGPU
+try:
+    GPU.initialize()
+except Exception as e:
+    print(f"ZeroGPU initialization failed: {e}")
+
+# Global data dictionary
 global_data = {
     'models': {},
     'tokens': {
@@ -57,7 +63,7 @@ class ModelManager:
             return {"model": Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename']), "name": model_config['name']}
         except Exception as e:
             print(f"Error loading model {model_config['name']}: {e}")
-            pass
+            return None
 
     def load_all_models(self):
         if self.loaded:
@@ -77,7 +83,6 @@ class ModelManager:
             return global_data['models']
         except Exception as e:
             print(f"Error loading models: {e}")
-            pass
             return {}
 
 model_manager = ModelManager()
@@ -112,28 +117,24 @@ def remove_repetitive_responses(responses):
         normalized_response = remove_duplicates(response['response'])
         if normalized_response not in seen:
             seen.add(normalized_response)
-
-
             unique_responses.append({'model': response['model'], 'response': normalized_response})
     return unique_responses
 
-@app.post("/chat/")
-@spaces.GPU(duration=0)
-async def chat(request: ChatRequest):
+@app.post("/generate/")
+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 = []
-            for future, model_name in zip(as_completed(futures), global_data['models'].keys()):
-                response = future.result()
-                responses.append({'model': model_name, 'response': response})
+            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"An error occurred: {e}")
+        raise HTTPException(status_code=500, detail=f"Error generating responses: {e}")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)