Update app.py

app.py

@@ -5,13 +5,18 @@ 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()
 
-# Global data storage
+try:
+    ZeroGPU.initialize()
+except Exception:
+    pass
+
 global_data = {
     'models': {},
     'tokens': {
@@ -55,8 +60,7 @@ 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 as e:
-            print(f"Error loading model {model_config['name']}: {e}")
+        except Exception:
             pass
 
     def load_all_models(self):
@@ -72,13 +76,12 @@ class ModelManager:
                     if model:
                         models.append(model)
             
-            global_data['models'] = models
+            global_data['models'] = {model['name']: model['model'] for model in models}
             self.loaded = True
-            return models
-        except Exception as e:
-            print(f"Error loading models: {e}")
+            return global_data['models']
+        except Exception:
             pass
-            return []
+            return {}
 
 model_manager = ModelManager()
 model_manager.load_all_models()
@@ -112,48 +115,30 @@ def remove_repetitive_responses(responses):
         normalized_response = remove_duplicates(response['response'])
         if normalized_response not in seen:
             seen.add(normalized_response)
-            unique_responses.append(response)
+            unique_responses.append({'model': response['model'], 'response': normalized_response})
     return unique_responses
 
+@app.post("/generate/")
 @spaces.GPU(duration=0)
-async def generate_model_response(model, inputs, top_k, top_p, temperature):
-    try:
-        responses = model.generate(inputs, top_k=top_k, top_p=top_p, temperature=temperature)
-        return responses
-    except Exception as e:
-        print(f"Error generating model response: {e}")
-        pass
-        return []
-
-@app.post("/generate")
 async def generate(request: ChatRequest):
     try:
-        if not global_data['models']:
-            raise HTTPException(status_code=500, detail="Models not loaded")
-
-        model = global_data['models'][0]['model']
-        inputs = normalize_input(request.message)
-        responses = await generate_model_response(model, inputs, request.top_k, request.top_p, request.temperature)
-        best_response = responses[0] if responses else {}
-        unique_responses = remove_repetitive_responses(responses)
-        return {
-            "best_response": best_response,
-            "all_responses": unique_responses
-        }
-    except Exception as e:
-        print(f"Error in generate endpoint: {e}")
-        pass
-        return {"error": str(e)}
-
-@app.api_route("/{method_name:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
-async def handle_request(method_name: str, request: Request):
-    try:
-        body = await request.json()
-        return {"message": "Request handled successfully", "body": body}
+        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']):
+                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}")
     except Exception as e:
-        print(f"Error handling request: {e}")
-        pass
-        return {"error": str(e)}
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)