import os import gc import io from llama_cpp import Llama from concurrent.futures import ThreadPoolExecutor, as_completed from fastapi import FastAPI, Request, HTTPException from fastapi.responses import JSONResponse from tqdm import tqdm from dotenv import load_dotenv from pydantic import BaseModel from huggingface_hub import hf_hub_download, login import spacy from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity import uvicorn import psutil import torch load_dotenv() app = FastAPI() HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") if HUGGINGFACE_TOKEN: login(token=HUGGINGFACE_TOKEN) global_data = { 'model_configs' == [ {"repo_id": "Ffftdtd5dtft/gpt2-xl-Q2_K-GGUF", "filename": "gpt2-xl-q2_k.gguf", "name": "GPT-2 XL"}, {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-instruct-q2_k.gguf", "name": "Meta Llama 3.1-8B Instruct"}, {"repo_id": "Ffftdtd5dtft/gemma-2-9b-it-Q2_K-GGUF", "filename": "gemma-2-9b-it-q2_k.gguf", "name": "Gemma 2-9B IT"}, {"repo_id": "Ffftdtd5dtft/gemma-2-27b-Q2_K-GGUF", "filename": "gemma-2-27b-q2_k.gguf", "name": "Gemma 2-27B"}, {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-Q2_K-GGUF", "filename": "phi-3-mini-128k-instruct-q2_k.gguf", "name": "Phi-3 Mini 128K Instruct"}, {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-8B-Q2_K-GGUF", "filename": "meta-llama-3.1-8b-q2_k.gguf", "name": "Meta Llama 3.1-8B"}, {"repo_id": "Ffftdtd5dtft/Qwen2-7B-Instruct-Q2_K-GGUF", "filename": "qwen2-7b-instruct-q2_k.gguf", "name": "Qwen2 7B Instruct"}, {"repo_id": "Ffftdtd5dtft/starcoder2-3b-Q2_K-GGUF", "filename": "starcoder2-3b-q2_k.gguf", "name": "Starcoder2 3B"}, {"repo_id": "Ffftdtd5dtft/Qwen2-1.5B-Instruct-Q2_K-GGUF", "filename": "qwen2-1.5b-instruct-q2_k.gguf", "name": "Qwen2 1.5B Instruct"}, {"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/Mistral-Nemo-Instruct-2407-Q2_K-GGUF", "filename": "mistral-nemo-instruct-2407-q2_k.gguf", "name": "Mistral Nemo Instruct 2407"}, {"repo_id": "Ffftdtd5dtft/Hermes-3-Llama-3.1-8B-IQ1_S-GGUF", "filename": "hermes-3-llama-3.1-8b-iq1_s-imat.gguf", "name": "Hermes 3 Llama 3.1-8B"}, {"repo_id": "Ffftdtd5dtft/Phi-3.5-mini-instruct-Q2_K-GGUF", "filename": "phi-3.5-mini-instruct-q2_k.gguf", "name": "Phi 3.5 Mini Instruct"}, {"repo_id": "Ffftdtd5dtft/Meta-Llama-3.1-70B-Instruct-Q2_K-GGUF", "filename": "meta-llama-3.1-70b-instruct-q2_k.gguf", "name": "Meta Llama 3.1-70B Instruct"}, {"repo_id": "Ffftdtd5dtft/codegemma-2b-IQ1_S-GGUF", "filename": "codegemma-2b-iq1_s-imat.gguf", "name": "Codegemma 2B"}, {"repo_id": "Ffftdtd5dtft/Phi-3-mini-128k-instruct-IQ2_XXS-GGUF", "filename": "phi-3-mini-128k-instruct-iq2_xxs-imat.gguf", "name": "Phi 3 Mini 128K Instruct XXS"}, {"repo_id": "Ffftdtd5dtft/TinyLlama-1.1B-Chat-v1.0-IQ1_S-GGUF", "filename": "tinyllama-1.1b-chat-v1.0-iq1_s-imat.gguf", "name": "TinyLlama 1.1B Chat"}, {"repo_id": "Ffftdtd5dtft/Mistral-NeMo-Minitron-8B-Base-IQ1_S-GGUF", "filename": "mistral-nemo-minitron-8b-base-iq1_s-imat.gguf", "name": "Mistral NeMo Minitron 8B Base"}, {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Instruct-2407-Q2_K-GGUF", "filename": "mistral-nemo-instruct-2407-q2_k.gguf", "name": "Mistral Nemo Instruct 2407"}, ] } class ModelManager: def __init__(self): self.models = {} self.load_models() def load_models(self): for config in tqdm(global_data['model_configs'], desc="Loading models"): model_name = config['name'] if model_name not in self.models: try: model_path = hf_hub_download(repo_id=config['repo_id'], use_auth_token=HUGGINGFACE_TOKEN) model = Llama.from_file(model_path, n_ctx=512, n_gpu=1) self.models[model_name] = model except Exception as e: self.models[model_name] = None finally: gc.collect() def get_model(self, model_name: str): return self.models.get(model_name) model_manager = ModelManager() class ChatRequest(BaseModel): message: str async def generate_model_response(model, inputs: str) -> str: try: if model: response = model(inputs, max_tokens=150) return response['choices'][0]['text'].strip() else: return "Model not loaded" except Exception as e: return f"Error: Could not generate a response. Details: {e}" async def process_message(message: str) -> dict: inputs = message.strip() responses = {} with ThreadPoolExecutor(max_workers=min(len(global_data['model_configs']), 4)) as executor: futures = [executor.submit(generate_model_response, model_manager.get_model(config['name']), inputs) for config in global_data['model_configs'] if model_manager.get_model(config['name'])] for i, future in enumerate(tqdm(as_completed(futures), total=len(futures), desc="Generating responses")): try: model_name = global_data['model_configs'][i]['name'] responses[model_name] = future.result() except Exception as e: responses[model_name] = f"Error processing {model_name}: {e}" nlp = spacy.load("en_core_web_sm") stop_words = spacy.lang.en.stop_words.STOP_WORDS def custom_tokenizer(text): doc = nlp(text) return [token.lemma_.lower() for token in doc if not token.is_stop and not token.is_punct] vectorizer = TfidfVectorizer(tokenizer=custom_tokenizer) reference_text = message response_texts = list(responses.values()) tfidf_matrix = vectorizer.fit_transform([reference_text] + response_texts) similarities = cosine_similarity(tfidf_matrix[0:1], tfidf_matrix[1:]) best_response_index = similarities.argmax() best_response_model = list(responses.keys())[best_response_index] best_response_text = response_texts[best_response_index] return {"best_response": {"model": best_response_model, "text": best_response_text}, "all_responses": responses} @app.post("/generate_multimodel") async def api_generate_multimodel(request: Request): try: data = await request.json() message = data.get("message") if not message: raise HTTPException(status_code=400, detail="Missing message") response = await process_message(message) return JSONResponse(response) except HTTPException as e: raise e except Exception as e: return JSONResponse({"error": str(e)}, status_code=500) @app.on_event("startup") async def startup_event(): pass @app.on_event("shutdown") async def shutdown_event(): gc.collect() def release_resources(): try: torch.cuda.empty_cache() gc.collect() except Exception as e: pass def resource_manager(): MAX_RAM_PERCENT = 20 MAX_CPU_PERCENT = 20 MAX_GPU_PERCENT = 20 MAX_RAM_MB = 2048 while True: try: virtual_mem = psutil.virtual_memory() current_ram_percent = virtual_mem.percent current_ram_mb = virtual_mem.used / (1024 * 1024) if current_ram_percent > MAX_RAM_PERCENT or current_ram_mb > MAX_RAM_MB: release_resources() current_cpu_percent = psutil.cpu_percent() if current_cpu_percent > MAX_CPU_PERCENT: psutil.Process(os.getpid()).nice() if torch.cuda.is_available(): gpu = torch.cuda.current_device() gpu_mem = torch.cuda.memory_percent(gpu) if gpu_mem > MAX_GPU_PERCENT: release_resources() except Exception as e: pass if __name__ == "__main__": import threading resource_thread = threading.Thread(target=resource_manager) resource_thread.daemon = True resource_thread.start() port = int(os.environ.get("PORT", 7860)) uvicorn.run(app, host="0.0.0.0", port=port)