import asyncio import os import torch import redis import io import numpy as np import base64 from PIL import Image from dotenv import load_dotenv from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments from transformers import pipeline from fastapi import FastAPI, HTTPException from fastapi.responses import HTMLResponse from typing import List, Dict, Any import logging # Configuración de logging logging.basicConfig(level=logging.DEBUG) logger = logging.getLogger(__name__) # Cargar variables de entorno load_dotenv() HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN") REDIS_HOST = os.getenv("REDIS_HOST") REDIS_PORT = os.getenv("REDIS_PORT") REDIS_PASSWORD = os.getenv("REDIS_PASSWORD") # Configuración de Redis redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True) # Inicializar la aplicación FastAPI app = FastAPI() # Diccionario para almacenar los modelos y sus propiedades en memoria model_dict: Dict[str, Any] = {} model_properties: Dict[str, Dict[str, Any]] = {} model_lock = asyncio.Lock() # Lista para almacenar el historial de mensajes y datos de entrenamiento en memoria message_history: List[str] = [] TRAINING_DATA: List[Dict[str, torch.Tensor]] = [] # Datos de entrenamiento para texto MUSIC_TRAINING_DATA: List[Dict[str, torch.Tensor]] = [] # Datos de entrenamiento para música IMAGE_TRAINING_DATA: List[Dict[str, torch.Tensor]] = [] # Datos de entrenamiento para imágenes # Parámetros adicionales para la generación de respuestas TEMPERATURE = 0.7 TOP_PROBABILITY = 0.9 TOP_K = 50 FREQUENCY_PENALTY = 0.7 MAX_TOKENS = 1024 # Límite máximo de tokens por respuesta UNIQUE_RESPONSES = set() # Conjunto para almacenar respuestas únicas # Inicializar el pipeline de generación de música y de imágenes musicgen_pipeline = pipeline("text-to-audio", model="facebook/musicgen-small") image_pipeline = pipeline("text-to-image", model="black-forest-labs/FLUX.1-schnell") # Función para almacenar en Redis def store_in_redis(key: str, value: Any): if isinstance(value, bytes): redis_client.set(key, value) else: redis_client.set(key, str(value)) # Función para recuperar de Redis def retrieve_from_redis(key: str): value = redis_client.get(key) if value is None: return None try: return value.encode('latin1') # Decodificar si es bytes except AttributeError: return value # Si es texto # Cargar modelos sincrónicamente al iniciar la aplicación async def load_models(): global model_dict, model_properties if not model_dict: # Solo cargar si el diccionario está vacío for model_name in ["gpt2-medium", "gpt2-large", "gpt2", "google/gemma-2-9b", "meta-llama/Meta-Llama-3.1-8B-Instruct"]: model_key = f"model:{model_name}" tokenizer_key = f"tokenizer:{model_name}" model_data = retrieve_from_redis(model_key) tokenizer_data = retrieve_from_redis(tokenizer_key) if model_data and tokenizer_data: model = torch.load(io.BytesIO(model_data)) tokenizer = torch.load(io.BytesIO(tokenizer_data)) model_dict[model_name] = (model, tokenizer) logger.info(f"Loaded {model_name} from Redis") else: model = AutoModelForCausalLM.from_pretrained(model_name) tokenizer = AutoTokenizer.from_pretrained(model_name) model_dict[model_name] = (model, tokenizer) # Guardar modelos y tokenizers en Redis store_in_redis(model_key, torch.save(model, io.BytesIO())) store_in_redis(tokenizer_key, torch.save(tokenizer, io.BytesIO())) model_properties[model_name] = { 'pad_token': tokenizer.pad_token, 'pad_token_id': tokenizer.pad_token_id, 'eos_token': tokenizer.eos_token, 'eos_token_id': tokenizer.eos_token_id, 'bos_token': tokenizer.bos_token, 'bos_token_id': tokenizer.bos_token_id, 'unk_token': tokenizer.unk_token, 'unk_token_id': tokenizer.unk_token_id, 'padding_side': tokenizer.padding_side, 'special_tokens_map': tokenizer.special_tokens_map, 'model': model, 'tokenizer': tokenizer } logger.info(f"Successfully loaded {model_name} model and tokenizer") # Cargar modelos una vez al iniciar la aplicación asyncio.run(load_models()) # Funciones para las APIs adicionales de música e imágenes def generate_music(prompt: str) -> bytes: # Generación de música utilizando el pipeline en memoria audio = musicgen_pipeline(prompt)['audio'] return audio def generate_image(prompt: str) -> bytes: # Generación de imagen utilizando el pipeline en memoria outputs = image_pipeline(prompt)["sample"][0] buffered = io.BytesIO() outputs.save(buffered, format="PNG") return buffered.getvalue() # Ruta principal para la interfaz web @app.get('/') async def main(): html_code = """ ChatGPT Chatbot
""" return HTMLResponse(content=html_code) # Ruta para generar contenido basado en la consulta @app.post('/generate') async def generate_content(query: str): async def generate_unique_response(q): attempts = 0 while attempts < 5: responses = await generate_responses(q) unique_responses = [response for response in responses if is_unique(response)] if unique_responses: parts = [] for response in unique_responses: parts.extend(split_response(response, model_properties[next(iter(model_dict))]['tokenizer'])) return parts attempts += 1 raise HTTPException(status_code=500, detail="No unique response found after multiple attempts") def is_unique(response): if response in UNIQUE_RESPONSES: return False else: UNIQUE_RESPONSES.add(response) return True async def generate_responses(q): responses = [] for model_name, (model, tokenizer) in model_dict.items(): input_ids = tokenizer.encode(q, return_tensors='pt') output = model.generate( input_ids, max_length=MAX_TOKENS, num_return_sequences=1, temperature=TEMPERATURE, top_p=TOP_PROBABILITY, top_k=TOP_K, frequency_penalty=FREQUENCY_PENALTY ) response = tokenizer.decode(output[0], skip_special_tokens=True) responses.append(response) return responses async def train_model(): global TRAINING_DATA if not TRAINING_DATA: raise ValueError("No training data available") model_name = 'gpt2' # Nombre del modelo a usar para entrenamiento tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForCausalLM.from_pretrained(model_name) training_args = TrainingArguments( output_dir='./results', # Directorio de resultados (en memoria, se puede ignorar) per_device_train_batch_size=4, num_train_epochs=1, save_steps=10_000, save_total_limit=2, ) trainer = Trainer( model=model, args=training_args, train_dataset=TRAINING_DATA ) trainer.train() # Guardar el modelo entrenado en Redis model_key = "model:trained" tokenizer_key = "tokenizer:trained" store_in_redis(model_key, torch.save(model, io.BytesIO())) store_in_redis(tokenizer_key, torch.save(tokenizer, io.BytesIO())) return model, tokenizer async def auto_learn(): global TRAINING_DATA if message_history: new_data = "\n".join(message_history) TRAINING_DATA.append(new_data) await train_model() async def auto_learn_music(): global MUSIC_TRAINING_DATA if MUSIC_TRAINING_DATA: inputs = musicgen_pipeline.tokenizer(MUSIC_TRAINING_DATA, return_tensors="pt", padding=True) model = musicgen_pipeline.model model.train() optimizer = torch.optim.Adam(model.parameters(), lr=5e-5) loss_fn = torch.nn.CrossEntropyLoss() for epoch in range(1): outputs = model(**inputs) loss = loss_fn(outputs.logits, inputs['labels']) optimizer.zero_grad() loss.backward() optimizer.step() MUSIC_TRAINING_DATA = [] async def auto_learn_images(): global IMAGE_TRAINING_DATA if IMAGE_TRAINING_DATA: for image_data in IMAGE_TRAINING_DATA: image = Image.open(io.BytesIO(image_data)) image_tensor = torch.tensor(np.array(image)).unsqueeze(0) # Adaptar según el modelo # Implementar el entrenamiento del modelo aquí con `image_tensor` model = image_pipeline.model model.train() optimizer = torch.optim.Adam(model.parameters(), lr=1e-5) loss_fn = torch.nn.MSELoss() target_tensor = torch.zeros_like(image_tensor) # Definir `target_tensor` según sea necesario for epoch in range(1): outputs = model(image_tensor) loss = loss_fn(outputs, target_tensor) optimizer.zero_grad() loss.backward() optimizer.step() IMAGE_TRAINING_DATA = [] def generate_music_from_api(prompt: str) -> bytes: # Llamada a la API para generar música audio = generate_music(prompt) store_in_redis(f"music:{prompt}", audio) # Almacenar música en Redis return audio def generate_image_from_api(prompt: str) -> bytes: # Llamada a la API para generar imágenes image = generate_image(prompt) store_in_redis(f"image:{prompt}", image) # Almacenar imagen en Redis return image try: tokenizer = model_properties[next(iter(model_dict))]['tokenizer'] final_responses = await generate_unique_response(query) await auto_learn() await auto_learn_music() await auto_learn_images() music = generate_music_from_api(query) image = generate_image_from_api(query) # Convertir la imagen a base64 para mostrarla en la interfaz web buffered = io.BytesIO(image) img_str = base64.b64encode(buffered.getvalue()).decode('utf-8') return {"responses": final_responses, "music": base64.b64encode(music).decode('utf-8'), "image": img_str} except Exception as e: logger.error(f"Error processing the request: {e}") raise HTTPException(status_code=500, detail="Error processing the request") # Ruta para generación de música @app.post('/music') async def generate_music_endpoint(prompt: str): try: music = generate_music_from_api(prompt) return {"music": base64.b64encode(music).decode('utf-8')} except Exception as e: logger.error(f"Error generating music: {e}") raise HTTPException(status_code=500, detail="Error generating music") # Ruta para generación de imágenes @app.post('/image') async def generate_image_endpoint(prompt: str): try: image = generate_image_from_api(prompt) img_str = base64.b64encode(image).decode('utf-8') return {"image": img_str} except Exception as e: logger.error(f"Error generating image: {e}") raise HTTPException(status_code=500, detail="Error generating image") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)