Update app.py

app.py

@@ -2,7 +2,7 @@ from llama_cpp import Llama
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import re
 import uvicorn
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 import os
 from dotenv import load_dotenv
@@ -16,9 +16,7 @@ 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'}}
 
-model_configs = [
-    {"repo_id": "Hjgugugjhuhjggg/mergekit-ties-tzamfyy-Q2_K-GGUF", "filename": "mergekit-ties-tzamfyy-q2_k.gguf", "name": "my_model"}
-]
+model_configs = [{"repo_id": "Hjgugugjhuhjggg/mergekit-ties-tzamfyy-Q2_K-GGUF", "filename": "mergekit-ties-tzamfyy-q2_k.gguf", "name": "my_model"}]
 
 models = {}
 
@@ -40,7 +38,6 @@ for config in model_configs:
         print(f"Failed to load model {config['name']}. Exiting.")
         exit(1)
 
-
 class ChatRequest(BaseModel):
     message: str
 
@@ -62,7 +59,7 @@ def generate_model_response(model, inputs):
     try:
         if model is None:
             return "Model loading failed."
-        response = model(inputs)
+        response = model(inputs, max_tokens=-1)
         return remove_duplicates(response['choices'][0]['text'])
     except Exception as e:
         print(f"Error generating response: {e}")
@@ -81,20 +78,25 @@ app.add_middleware(
 @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 models.values()]
-        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']
-
-    formatted_response = ""
-    for model, response in unique_responses.items():
-        formatted_response += f"**{model}:**\n{response}\n\n"
-
-    return {"response": formatted_response}
+    chunk_size = 500
+    chunks = [inputs[i:i + chunk_size] for i in range(0, len(inputs), chunk_size)]
+    overall_response = ""
+    for chunk in chunks:
+        with ThreadPoolExecutor() as executor:
+            futures = [executor.submit(generate_model_response, model, chunk) 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 in responses:
+            if response['model'] not in unique_responses and response['response']:
+                unique_responses[response['model']] = response['response']
+
+        chunk_response = ""
+        for model, response in unique_responses.items():
+            chunk_response += f"**{model}:**\n{response}\n\n"
+        overall_response += chunk_response
+
+    return {"response": overall_response}
 
 async def process_message(message, history):
     try:
@@ -106,12 +108,10 @@ async def process_message(message, history):
     except requests.exceptions.RequestException as e:
         return history, f"Error communicating with the backend: {e}"
 
+
 iface = gr.Interface(
     fn=process_message,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Enter your message here..."),
-        gr.State([])
-    ],
+    inputs=[gr.Textbox(lines=2, placeholder="Enter your message here..."), gr.State([])],
     outputs=[gr.Chatbot(), gr.State([])],
     title="Multi-Model LLM API",
     description="Enter a message and get responses from multiple LLMs.",