Update app.py

app.py

@@ -1,112 +1,70 @@
 from pydantic import BaseModel
 from llama_cpp import Llama
-import re
 import os
 import gradio as gr
 from dotenv import load_dotenv
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
 import spaces
-import urllib3
+import asyncio
 import random
 
-urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
-
 app = FastAPI()
 load_dotenv()
 
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
-global_data = {
-    'model': None,
-    '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'
-    }
-}
-
-model_configs = [
-    {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/codegemma-2b-IQ1_S-GGUF", "filename": "codegemma-2b-iq1_s-imat.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Phi-3.5-mini-instruct-Q2_K-GGUF", "filename": "phi-3.5-mini-instruct-q2_k.gguf"},
-    {"repo_id": "Ffftdtd5dtft/TinyLlama-1.1B-Chat-v1.0-IQ1_S-GGUF", "filename": "tinyllama-1.1b-chat-v1.0-iq1_s-imat.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Minitron-8B-Base-IQ1_S-GGUF", "filename": "mistral-nemo-minitron-8b-base-iq1_s-imat.gguf"},
-    {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Instruct-2407-Q2_K-GGUF", "filename": "mistral-nemo-instruct-2407-q2_k.gguf"}
-]
-
 class ModelManager:
     def __init__(self):
-        self.model = None
+        self.model = self.load_models()
 
     def load_models(self):
         models = []
+        model_configs = [
+            {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/codegemma-2b-IQ1_S-GGUF", "filename": "codegemma-2b-iq1_s.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Phi-3.5-mini-instruct-Q2_K-GGUF", "filename": "phi-3.5-mini-instruct-q2_k.gguf"},
+            {"repo_id": "Ffftdtd5dtft/TinyLlama-1.1B-Chat-v1.0-IQ1_S-GGUF", "filename": "tinyllama-1.1b-chat-v1.0-iq1_s.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Minitron-8B-Base-IQ1_S-GGUF", "filename": "mistral-nemo-minitron-8b-base-iq1_s.gguf"},
+            {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Instruct-2407-Q2_K-GGUF", "filename": "mistral-nemo-instruct-2407-q2_k.gguf"}
+        ]
         for config in model_configs:
-            try:
-                model = Llama.from_pretrained(repo_id=config['repo_id'], filename=config['filename'], use_auth_token=HUGGINGFACE_TOKEN)
-                models.append(model)
-            except Exception:
-                continue
-        self.model = models
+            model = Llama.from_pretrained(repo_id=config['repo_id'], filename=config['filename'], use_auth_token=HUGGINGFACE_TOKEN)
+            models.append(model)
+        return models
 
 model_manager = ModelManager()
-model_manager.load_models()
-global_data['model'] = model_manager.model
 
 class ChatRequest(BaseModel):
     message: str
 
-def normalize_input(input_text):
-    return input_text.strip()
-
-def remove_duplicates(text):
-    text = re.sub(r'(Hello there, how are you\? \[/INST\]){2,}', 'Hello there, how are you? [/INST]', text)
-    text = re.sub(r'(How are you\? \[/INST\]){2,}', 'How are you? [/INST]', text)
-    text = text.replace('[/INST]', '')
-    lines = text.split('\n')
-    unique_lines = []
-    seen_lines = set()
-    for line in lines:
-        if line not in seen_lines:
-            unique_lines.append(line)
-            seen_lines.add(line)
-    return '\n'.join(unique_lines)
-
 @spaces.GPU()
 async def generate_combined_response(inputs):
     combined_response = ""
     top_p = round(random.uniform(0.01, 1.00), 2)
     top_k = random.randint(1, 100)
     temperature = round(random.uniform(0.01, 2.00), 2)
-    for model in global_data['model']:
-        try:
-            response = model(inputs, top_p=top_p, top_k=top_k, temperature=temperature)
-            combined_response += remove_duplicates(response['choices'][0]['text']) + "\n"
-        except Exception:
-            continue
+    tasks = []
+    for model in model_manager.model:
+        tasks.append(model(inputs, top_p=top_p, top_k=top_k, temperature=temperature))
+    responses = await asyncio.gather(*tasks)
+    for response in responses:
+        combined_response += response['choices'][0]['text'] + "\n"
     return combined_response
 
 async def process_message(message):
-    inputs = normalize_input(message)
+    inputs = message.strip()
     combined_response = await generate_combined_response(inputs)
-    formatted_response = ""
-    for line in combined_response.split("\n"):
-        formatted_response += f"{line}\n\n"
-    return formatted_response
+    return combined_response
 
 @app.post("/generate_multimodel")
 async def api_generate_multimodel(request: Request):
@@ -119,10 +77,9 @@ iface = gr.Interface(
     fn=process_message,
     inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
     outputs=gr.Markdown(),
-    title="Multi-Model LLM API",
-    description="Enter a message and get responses from a unified model.",
+    title="Unified Multi-Model API",
+    description="Enter a message to get responses from a unified model."
 )
 
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", 7860))
-    iface.launch(server_port=port)
+    iface.launch()