Update app.py

app.py

@@ -4,10 +4,7 @@ from llama_cpp import Llama
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import uvicorn
 import re
-from dotenv import load_dotenv
-import spaces
-
-load_dotenv()
+from spaces import GPU
 
 app = FastAPI()
 
@@ -50,40 +47,29 @@ model_configs = [
 class ModelManager:
     def __init__(self):
         self.loaded = False
+        self.models = {}
 
     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:
-            pass
+        if model_config['name'] not in self.models:
+            try:
+                self.models[model_config['name']] = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'])
+            except Exception as e:
+                print(f"Error loading model {model_config['name']}: {e}")
 
     def load_all_models(self):
-        if self.loaded:
-            return global_data['models']
-        
-        try:
+        if not self.loaded:
             with ThreadPoolExecutor() as executor:
-                futures = [executor.submit(self.load_model, config) for config in model_configs]
-                models = []
-                for future in as_completed(futures):
-                    model = future.result()
-                    if model:
-                        models.append(model)
-            
-            global_data['models'] = models
+                for config in model_configs:
+                    executor.submit(self.load_model, config)
             self.loaded = True
-            return models
-        except Exception:
-            pass
+
+        return self.models
 
 model_manager = ModelManager()
-model_manager.load_all_models()
+global_data['models'] = model_manager.load_all_models()
 
 class ChatRequest(BaseModel):
     message: str
-    top_k: int = 50
-    top_p: float = 0.95
-    temperature: float = 0.7
 
 def normalize_input(input_text):
     return input_text.strip()
@@ -97,61 +83,50 @@ def remove_duplicates(text):
     seen_lines = set()
     for line in lines:
         if line not in seen_lines:
-            seen_lines.add(line)
             unique_lines.append(line)
+            seen_lines.add(line)
     return '\n'.join(unique_lines)
 
-def remove_repetitive_responses(responses):
-    seen = set()
-    unique_responses = []
-    for response in responses:
-        normalized_response = remove_duplicates(response['response'])
-        if normalized_response not in seen:
-            seen.add(normalized_response)
-            unique_responses.append(response)
-    return unique_responses
-
-def generate_chat_response(request, model_data):
-    model = model_data['model']
+@GPU(duration=0)
+def generate_model_response(model, inputs):
     try:
-        user_input = normalize_input(request.message)
-        response = model(user_input, top_k=request.top_k, top_p=request.top_p, temperature=request.temperature)
-        return response
-    except Exception:
-        pass
+        response = model(inputs)
+        return remove_duplicates(response['choices'][0]['text'])
+    except Exception as e:
+        print(f"Error generating model response: {e}")
+        return ""
 
 @app.post("/generate")
 async def generate(request: ChatRequest):
     try:
-        responses = []
-        models = global_data['models']
-        for model_data in models:
-            response = generate_chat_response(request, model_data)
-            if response:
-                responses.append({
-                    "model": model_data['name'],
-                    "response": response
-                })
-        
-        if not responses:
-            raise HTTPException(status_code=500, detail="Error: No responses generated.")
-        
-        responses = remove_repetitive_responses(responses)
-        best_response = responses[0] if responses else {}
-        return {
-            "best_response": best_response,
-            "all_responses": responses
-        }
-    except Exception:
-        pass
-
-@app.api_route("/{method_name:path}", methods=["GET", "POST", "PUT", "DELETE", "PATCH"])
-async def handle_request(method_name: str, request: Request):
+        inputs = normalize_input(request.message)
+        with ThreadPoolExecutor() as executor:
+            futures = [
+                executor.submit(generate_model_response, model, inputs)
+                for model in global_data['models'].values()
+            ]
+            responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(global_data['models'].keys(), as_completed(futures))]
+        unique_responses = remove_repetitive_responses(responses)
+        return unique_responses
+    except Exception as 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:
-        body = await request.json()
-        return {"message": "Request handled successfully", "body": body}
-    except Exception:
-        pass
+        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")
+
+def remove_repetitive_responses(responses):
+    unique_responses = {}
+    for response in responses:
+        if response['model'] not in unique_responses:
+            unique_responses[response['model']] = response['response']
+    return unique_responses
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)