Update app.py

app.py

@@ -1,3 +1,4 @@
+
 from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
 from llama_cpp import Llama
@@ -5,18 +6,12 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 import uvicorn
 import re
 from dotenv import load_dotenv
-from spaces.zero import ZeroGPU
 import spaces
 
 load_dotenv()
 
 app = FastAPI()
 
-try:
-    ZeroGPU.initialize()
-except Exception:
-    pass
-
 global_data = {
     'models': {},
     'tokens': {
@@ -60,7 +55,8 @@ class ModelManager:
     def load_model(self, model_config):
         try:
             return {"model": Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename']), "name": model_config['name']}
-        except Exception:
+        except Exception as e:
+            print(f"Error loading model {model_config['name']}: {e}")
             pass
 
     def load_all_models(self):
@@ -79,7 +75,8 @@ class ModelManager:
             global_data['models'] = {model['name']: model['model'] for model in models}
             self.loaded = True
             return global_data['models']
-        except Exception:
+        except Exception as e:
+            print(f"Error loading models: {e}")
             pass
             return {}
 
@@ -115,12 +112,14 @@ 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("/generate/")
+@app.post("/chat/")
 @spaces.GPU(duration=0)
-async def generate(request: ChatRequest):
+async def chat(request: ChatRequest):
     try:
         normalized_message = normalize_input(request.message)
         with ThreadPoolExecutor() as executor:
@@ -128,17 +127,13 @@ async def generate(request: ChatRequest):
                                        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']):
-                generated_text = future.result()
-                responses.append({'model': model_name, 'response': generated_text})
-        
-        return remove_repetitive_responses(responses)
-    except NotImplementedError as nie:
-        raise HTTPException(status_code=500, detail=str(nie))
-    except ZeroGPU.ZeroGPUException as gpu_exc:
-        raise HTTPException(status_code=500, detail=f"ZeroGPU Error: {gpu_exc}")
+            for future, model_name in zip(as_completed(futures), global_data['models'].keys()):
+                response = future.result()
+                responses.append({'model': model_name, 'response': response})
+        unique_responses = remove_repetitive_responses(responses)
+        return unique_responses
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=f"An error occurred: {e}")
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8000)