Update apghp.py

apghp.py

@@ -1,26 +1,34 @@
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+import os
+import gc
+import io
 from llama_cpp import Llama
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from fastapi import FastAPI, Request, HTTPException, Lifespan
+from fastapi.responses import JSONResponse
 from tqdm import tqdm
-import uvicorn
 from dotenv import load_dotenv
-from difflib import SequenceMatcher
-import re
-import spaces  # Importar la librería spaces
+from pydantic import BaseModel
+from huggingface_hub import hf_hub_download, login
+from nltk.tokenize import word_tokenize
+from nltk.corpus import stopwords
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import nltk
+import uvicorn
+import psutil
+import torch
+import tempfile
+
+nltk.download('punkt')
+nltk.download('stopwords')
 
-# Cargar variables de entorno
 load_dotenv()
 
-# Inicializar aplicación FastAPI
 app = FastAPI()
+HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+if HUGGINGFACE_TOKEN:
+    login(token=HUGGINGFACE_TOKEN)
 
-# Diccionario global para almacenar los modelos
-global_data = {
-    'models': []
-}
-
-# Configuración de los modelos
 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"},
@@ -43,146 +51,139 @@ model_configs = [
     {"repo_id": "Ffftdtd5dtft/Mistral-Nemo-Instruct-2407-Q2_K-GGUF", "filename": "mistral-nemo-instruct-2407-q2_k.gguf", "name": "Mistral Nemo Instruct 2407"}
 ]
 
-# Clase para gestionar modelos
+global_data = {'model_configs': model_configs, 'training_data': io.StringIO()}
+
 class ModelManager:
     def __init__(self):
-        self.models = []
-        self.loaded = False  # Para verificar si ya están cargados
-    
-    def load_model(self, model_config):
-        print(f"Cargando modelo: {model_config['name']}...")
-        return {"model": Llama.from_pretrained(repo_id=model_config['repo_id'], filename=model_config['filename']), "name": model_config['name']}
-    
-    def load_all_models(self):
-        if self.loaded:  # Si los modelos ya están cargados, no los vuelve a cargar
-            print("Modelos ya están cargados. No es necesario volver a cargarlos.")
-            return self.models
-        
-        print("Iniciando carga de modelos...")
-        with ThreadPoolExecutor() as executor:  # No hay límite de trabajadores
-            futures = [executor.submit(self.load_model, config) for config in model_configs]
-            models = []
-            for future in tqdm(as_completed(futures), total=len(model_configs), desc="Cargando modelos", unit="modelo"):
+        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 = future.result()
-                    models.append(model)
-                    print(f"Modelo cargado exitosamente: {model['name']}")
+                    with tempfile.NamedTemporaryFile(suffix=".gguf", delete=False) as temp_file:
+                        model_path = hf_hub_download(repo_id=config['repo_id'], filename=temp_file.name, use_auth_token=HUGGINGFACE_TOKEN)
+                        model = Llama.from_file(model_path, n_ctx=512, n_gpu=1)
+                        self.models[model_name] = model
+                        print(f"Model '{model_name}' loaded successfully.")
+                        os.remove(temp_file.name) #remove the temp file after loading
                 except Exception as e:
-                    print(f"Error al cargar el modelo: {e}")
-        
-        self.models = models
-        self.loaded = True  # Marcar como cargados
-        print("Todos los modelos han sido cargados.")
-        return self.models
-
-# Instanciar ModelManager
-model_manager = ModelManager()
+                    print(f"Error loading model {model_name}: {e}")
+                    self.models[model_name] = None
+                finally:
+                    gc.collect()
 
-# Cargar modelos al iniciar la aplicación, solo la primera vez
-global_data['models'] = model_manager.load_all_models()
+    def get_model(self, model_name: str):
+        return self.models.get(model_name)
+
+
+model_manager = ModelManager()
 
-# Modelo global para la solicitud de chat
 class ChatRequest(BaseModel):
     message: str
-    top_k: int = 50
-    top_p: float = 0.95
-    temperature: float = 0.7
 
-# Función para generar respuestas de chat
-@spaces.GPU(duration=0)  # Anotación para usar GPU con duración 0
-def generate_chat_response(request, model_data):
+async def generate_model_response(model, inputs: str) -> str:
     try:
-        user_input = normalize_input(request.message)
-        llm = model_data['model']
-        response = llm.create_chat_completion(
-            messages=[{"role": "user", "content": user_input}],
-            top_k=request.top_k,
-            top_p=request.top_p,
-            temperature=request.temperature
-        )
-        reply = response['choices'][0]['message']['content']
-        return {"response": reply, "literal": user_input, "model_name": model_data['name']}
+        if model:
+            response = model(inputs, max_tokens=150)
+            return response['choices'][0]['text'].strip()
+        else:
+            return "Model not loaded"
     except Exception as e:
-        return {"response": f"Error: {str(e)}", "literal": user_input, "model_name": model_data['name']}
-
-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 = list(dict.fromkeys(lines))
-    return '\n'.join(unique_lines).strip()
-
-def remove_repetitive_responses(responses):
-    seen = set()
-    unique_responses = []
-    for response in responses:
-        normalized_response = remove_duplicates(response['response'])
-        if normalized_response not in seen:
-            seen.add(normalized_response)
-            unique_responses.append(response)
-    return unique_responses
-
-def select_best_response(responses):
-    print("Filtrando respuestas...")
-    responses = remove_repetitive_responses(responses)
-    responses = [remove_duplicates(response['response']) for response in responses]
-    unique_responses = list(set(responses))
-    coherent_responses = filter_by_coherence(unique_responses)
-    best_response = filter_by_similarity(coherent_responses)
-    return best_response
-
-def filter_by_coherence(responses):
-    print("Ordenando respuestas por coherencia...")
-    responses.sort(key=len, reverse=True)
-    return responses
-
-def filter_by_similarity(responses):
-    print("Filtrando respuestas por similitud...")
-    responses.sort(key=len, reverse=True)
-    best_response = responses[0]
-    for i in range(1, len(responses)):
-        ratio = SequenceMatcher(None, best_response, responses[i]).ratio()
-        if ratio < 0.9:
-            best_response = responses[i]
-            break
-    return best_response
-
-def worker_function(model_data, request):
-    print(f"Generando respuesta con el modelo: {model_data['name']}...")
-    response = generate_chat_response(request, model_data)
-    return response
-
-@app.post("/generate_chat")
-async def generate_chat(request: ChatRequest):
-    if not request.message.strip():
-        raise HTTPException(status_code=400, detail="The message cannot be empty.")
-    
-    print(f"Procesando solicitud: {request.message}")
-
-    responses = []
-    num_models = len(global_data['models'])
-
-    with ThreadPoolExecutor() as executor:  # No se establece límite de concurrencia
-        futures = [executor.submit(worker_function, model_data, request) for model_data in global_data['models']]
-        for future in tqdm(as_completed(futures), total=num_models, desc="Generando respuestas", unit="modelo"):
+        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:
-                response = future.result()
-                responses.append(response)
-            except Exception as exc:
-                print(f"Error en la generación de respuesta: {exc}")
+                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}"
+
+    stop_words = set(stopwords.words('english'))
+    vectorizer = TfidfVectorizer(tokenizer=word_tokenize, stop_words=stop_words)
+    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)
+
+
+async def startup():
+    pass
+
+async def shutdown():
+    gc.collect()
+
+app.add_event_handler("startup", startup)
+app.add_event_handler("shutdown", shutdown)
+
+def release_resources():
+    try:
+        torch.cuda.empty_cache()
+        gc.collect()
+    except Exception as e:
+        print(f"Failed to release resources: {e}")
+
+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)
 
-    best_response = select_best_response(responses)
-    
-    print(f"Mejor respuesta seleccionada: {best_response}")
+                if gpu_mem > MAX_GPU_PERCENT:
+                    release_resources()
 
-    return {
-        "best_response": best_response,
-        "all_responses": responses
-    }
+        except Exception as e:
+            print(f"Error in resource manager: {e}")
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    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)