Update app.py

app.py

@@ -8,6 +8,8 @@ from dotenv import load_dotenv
 from pydantic import BaseModel
 import requests
 import traceback
+from sentence_transformers import SentenceTransformer, util
+from gptcache import Cache, get_cache
 
 load_dotenv()
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
@@ -27,7 +29,9 @@ global_data = {
 }
 
 model_configs = [
-    {"repo_id": "Hjgugugjhuhjggg/mergekit-ties-tzamfyy-Q2_K-GGUF", "filename": "mergekit-ties-tzamfyy-q2_k.gguf", "name": "my_model"}
+    {"repo_id": "Hjgugugjhuhjggg/mergekit-ties-tzamfyy-Q2_K-GGUF", "filename": "mergekit-ties-tzamfyy-q2_k.gguf", "name": "my_model"},
+    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-q2_k.gguf", "name": "meta-llama-3.1-70b"},
+    {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf", "name": "gemma-2-27b"}
 ]
 
 models = {}
@@ -36,7 +40,7 @@ def load_model(model_config):
     model_name = model_config['name']
     if model_name not in models:
         try:
-            model = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'], use_auth_token=HUGGINGFACE_TOKEN)
+            model = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'], use_auth_token=HUGGINGFACE_TOKEN, verbose=True)
             models[model_name] = model
             global_data['models'] = models
             return model
@@ -80,7 +84,7 @@ def generate_model_response(model, inputs, max_tokens_per_part):
 
         text = response['choices'][0]['text']
         if text:
-          responses.append(remove_duplicates(text))
+            responses.append(remove_duplicates(text))
 
         return responses
 
@@ -89,6 +93,7 @@ def generate_model_response(model, inputs, max_tokens_per_part):
         traceback.print_exc()
         return [f"Error: {e}"]
 
+
 app = FastAPI()
 origins = ["*"]
 app.add_middleware(
@@ -99,28 +104,45 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+cache = get_cache()
+
 @app.post("/generate")
 async def generate(request: ChatRequest):
     inputs = normalize_input(request.message)
+    cached_result = cache.get(inputs)
+    if cached_result:
+        return {"response": cached_result}
+
     with ThreadPoolExecutor() as executor:
         futures = [executor.submit(generate_model_response, model, inputs, request.max_tokens_per_part) for model in models.values()]
         responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(models.keys(), as_completed(futures))]
 
-    unique_responses = {}
-    for response_set in responses:
-        model_name = response_set['model']
-        if model_name not in unique_responses:
-            unique_responses[model_name] = []
-        unique_responses[model_name].extend(response_set['response'])
+    model_embeddings = {}
+    sentence_model = SentenceTransformer('all-mpnet-base-v2')
+
+    all_responses = {}
+    for res in responses:
+        all_responses[res['model']] = " ".join(res['response'])
 
+    for model_name, response_text in all_responses.items():
+        embeddings = sentence_model.encode(response_text)
+        model_embeddings[model_name] = embeddings
 
-    formatted_response = ""
-    for model, response_parts in unique_responses.items():
-        formatted_response += f"**{model}:**\n"
-        for i, part in enumerate(response_parts):
-            formatted_response += f"Part {i+1}:\n{part}\n\n"
+    best_model = None
+    best_similarity = -1
 
+    for model_name1, embedding1 in model_embeddings.items():
+        avg_similarity = 0
+        for model_name2, embedding2 in model_embeddings.items():
+            cosine_scores = util.cos_sim(embedding1, embedding2)
+            avg_similarity += cosine_scores.item()
+        avg_similarity /= len(model_embeddings)
+        if avg_similarity > best_similarity:
+            best_similarity = avg_similarity
+            best_model = model_name1
 
+    formatted_response = f"**Best Model ({best_model}):**\n{all_responses[best_model]}\n\n"
+    cache.set(inputs, formatted_response)
     return {"response": formatted_response}
 
 if __name__ == "__main__":