from pydantic import BaseModel from llama_cpp import Llama import os import gradio as gr from dotenv import load_dotenv from fastapi import FastAPI, Request from fastapi.responses import JSONResponse import spaces import asyncio import random app = FastAPI() load_dotenv() HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") class ModelManager: def __init__(self): 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/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-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/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"} ] for config in model_configs: 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() class ChatRequest(BaseModel): message: str @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) 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 = message.strip() combined_response = await generate_combined_response(inputs) return combined_response @app.post("/generate_multimodel") async def api_generate_multimodel(request: Request): data = await request.json() message = data["message"] formatted_response = await process_message(message) return JSONResponse({"response": formatted_response}) iface = gr.Interface( fn=process_message, inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."), outputs=gr.Markdown(), title="Unified Multi-Model API", description="Enter a message to get responses from a unified model." ) if __name__ == "__main__": iface.launch()