Update app.py

app.py

@@ -2,9 +2,9 @@ from pydantic import BaseModel
 from llama_cpp import Llama
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import re
-import httpx
-import asyncio
-import gradio as gr
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 import os
 from dotenv import load_dotenv
 import json
@@ -13,50 +13,40 @@ load_dotenv()
 
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
-global_data = {
-    'models': {},
-    'tokens': {
-        'eos': 'eos_token',
-        'pad': 'pad_token',
-        'padding': 'padding_token',
-        'unk': 'unk_token',
-        'bos': 'bos_token',
-        'sep': 'sep_token',
-        'cls': 'cls_token',
-        'mask': 'mask_token'
-    }
-}
+app = FastAPI()
+
+origins = ["*"]  # Adjust as needed for production
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 
 model_configs = [
     {"repo_id": "Hjgugugjhuhjggg/mergekit-ties-tzamfyy-Q2_K-GGUF", "filename": "mergekit-ties-tzamfyy-q2_k.gguf", "name": "my_model"}
-    # Add more models here as needed.  Ensure the repo_id and filename are correct.
+    # Add more models here
 ]
 
-class ModelManager:
-    def __init__(self):
-        self.models = {}
+models = {}
+def load_model(model_config):
+    if model_config['name'] not in models:
+        try:
+            model = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'], use_auth_token=HUGGINGFACE_TOKEN)
+            models[model_config['name']] = model
+            print(f"Model '{model_config['name']}' loaded successfully.")
+            return model
+        except Exception as e:
+            print(f"Error loading model {model_config['name']}: {e}")
+            return None
 
-    def load_model(self, model_config):
-        if model_config['name'] not in self.models:
-            try:
-                model = Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename'], use_auth_token=HUGGINGFACE_TOKEN)
-                self.models[model_config['name']] = model
-                print(f"Model '{model_config['name']}' loaded successfully.")
-            except Exception as e:
-                print(f"Error loading model {model_config['name']}: {e}")
-                self.models[model_config['name']] = None # Indicate loading failure
+for config in model_configs:
+    load_model(config) #Load models on startup
 
 
-    def load_all_models(self):
-        with ThreadPoolExecutor() as executor:
-            futures = [executor.submit(self.load_model, config) for config in model_configs]
-            for future in as_completed(futures):
-                future.result() # Propagate exceptions during loading
-        return self.models
-
-model_manager = ModelManager()
-global_data['models'] = model_manager.load_all_models()
-
 
 class ChatRequest(BaseModel):
     message: str
@@ -69,7 +59,7 @@ def remove_duplicates(text):
     unique_lines = []
     seen_lines = set()
     for line in lines:
-        line = line.strip() #Added to remove extra whitespace
+        line = line.strip()
         if line and line not in seen_lines:
             unique_lines.append(line)
             seen_lines.add(line)
@@ -77,54 +67,37 @@ def remove_duplicates(text):
 
 def generate_model_response(model, inputs):
     try:
-        if model is None: #Handle failed model loading
+        if model is None:
             return ""
         response = model(inputs)
         return remove_duplicates(response['choices'][0]['text'])
     except Exception as e:
         print(f"Error generating model response: {e}")
-        return f"Error: {e}" #Return informative error message
-
-def remove_repetitive_responses(responses):
-    unique_responses = {}
-    for response in responses:
-        if response['model'] not in unique_responses and response['response']: #added check for empty responses
-            unique_responses[response['model']] = response['response']
-    return unique_responses
+        return f"Error: {e}"
 
 
-async def process_message(message, history):
-    inputs = normalize_input(message)
+@app.post("/generate")
+async def generate(request: ChatRequest):
+    inputs = normalize_input(request.message)
     with ThreadPoolExecutor() as executor:
         futures = [
             executor.submit(generate_model_response, model, inputs)
-            for model in global_data['models'].values()
+            for model in models.values()
         ]
-        responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(global_data['models'].keys(), as_completed(futures))]
+        responses = [{'model': model_name, 'response': future.result()} for model_name, future in zip(models.keys(), as_completed(futures))]
+
+    unique_responses = {}
+    for response in responses:
+        if response['model'] not in unique_responses and response['response']:
+            unique_responses[response['model']] = response['response']
 
-    unique_responses = remove_repetitive_responses(responses)
     formatted_response = ""
     for model, response in unique_responses.items():
         formatted_response += f"**{model}:**\n{response}\n\n"
 
-    history.append((message, formatted_response))
-    return history, ""
-
-
-iface = gr.Interface(
-    fn=process_message,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your message here..."),
-        gr.State([])
-    ],
-    outputs=[
-        gr.Chatbot(),
-        gr.Textbox(label="cURL command", visible=False) #Hidden cURL command
-    ],
-    title="Multi-Model LLM API",
-    description="Enter a message and get responses from multiple LLMs.",
-)
+    return {"response": formatted_response}
+
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))
-    iface.launch(server_port=port)
+    uvicorn.run(app, host="0.0.0.0", port=port)