|
from dotenv import load_dotenv |
|
import os |
|
import json |
|
import redis |
|
from transformers import ( |
|
AutoTokenizer, |
|
AutoModelForSequenceClassification, |
|
AutoModelForCausalLM, |
|
TrainingArguments, |
|
Trainer, |
|
AutoModelForTextToWaveform, |
|
pipeline, |
|
) |
|
from diffusers import FluxPipeline |
|
from fastapi import FastAPI, HTTPException, Request |
|
from fastapi.responses import HTMLResponse |
|
import multiprocessing |
|
import uuid |
|
import torch |
|
from torch.utils.data import Dataset |
|
import numpy as np |
|
|
|
load_dotenv() |
|
|
|
REDIS_HOST = os.getenv('REDIS_HOST') |
|
REDIS_PORT = os.getenv('REDIS_PORT') |
|
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD') |
|
|
|
app = FastAPI() |
|
|
|
default_language = "es" |
|
|
|
class ChatbotService: |
|
def __init__(self): |
|
self.redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
self.model_name = "response_model" |
|
self.tokenizer_name = "response_tokenizer" |
|
self.model = self.load_model_from_redis() |
|
self.tokenizer = self.load_tokenizer_from_redis() |
|
|
|
def get_response(self, user_id, message, language=default_language): |
|
if self.model is None or self.tokenizer is None: |
|
return "El modelo aún no está listo. Por favor, inténtelo de nuevo más tarde." |
|
input_text = f"Usuario: {message} Asistente:" |
|
input_ids = self.tokenizer.encode(input_text, return_tensors="pt").to("cpu") |
|
with torch.no_grad(): |
|
output = self.model.generate(input_ids=input_ids, max_length=100, num_beams=5, no_repeat_ngram_size=2, early_stopping=True) |
|
response = self.tokenizer.decode(output[0], skip_special_tokens=True) |
|
response = response.replace(input_text, "").strip() |
|
return response |
|
|
|
def load_model_from_redis(self): |
|
model_data_bytes = self.redis_client.get(f"model:{self.model_name}") |
|
if model_data_bytes: |
|
model = AutoModelForCausalLM.from_pretrained("gpt2") |
|
model.load_state_dict(torch.load(model_data_bytes)) |
|
return model |
|
return None |
|
|
|
def load_tokenizer_from_redis(self): |
|
tokenizer_data_bytes = self.redis_client.get(f"tokenizer:{self.tokenizer_name}") |
|
if tokenizer_data_bytes: |
|
tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
tokenizer.add_tokens(json.loads(tokenizer_data_bytes.decode("utf-8"))) |
|
tokenizer.pad_token = tokenizer.eos_token |
|
return tokenizer |
|
return None |
|
|
|
chatbot_service = ChatbotService() |
|
|
|
class UnifiedModel(AutoModelForSequenceClassification): |
|
def __init__(self, config): |
|
super().__init__(config) |
|
|
|
@staticmethod |
|
def load_model_from_redis(redis_client): |
|
model_name = "unified_model" |
|
model_path = f"models/{model_name}" |
|
if redis_client.exists(f"model:{model_name}"): |
|
redis_client.delete(f"model:{model_name}") |
|
if not os.path.exists(model_path): |
|
model = UnifiedModel.from_pretrained("gpt2", num_labels=3) |
|
model.save_pretrained(model_path) |
|
else: |
|
model = UnifiedModel.from_pretrained(model_path) |
|
return model |
|
|
|
class SyntheticDataset(Dataset): |
|
def __init__(self, tokenizer, data): |
|
self.tokenizer = tokenizer |
|
self.data = data |
|
|
|
def __len__(self): |
|
return len(self.data) |
|
|
|
def __getitem__(self, idx): |
|
item = self.data[idx] |
|
text = item['text'] |
|
label = item['label'] |
|
tokens = self.tokenizer(text, padding="max_length", truncation=True, max_length=128, return_tensors="pt") |
|
return {"input_ids": tokens["input_ids"].squeeze(), "attention_mask": tokens["attention_mask"].squeeze(), "labels": label} |
|
|
|
conversation_history = {} |
|
|
|
tokenizer_name = "unified_tokenizer" |
|
tokenizer = None |
|
unified_model = None |
|
musicgen_tokenizer = AutoTokenizer.from_pretrained("facebook/musicgen-small") |
|
musicgen_model = AutoModelForTextToWaveform.from_pretrained("facebook/musicgen-small") |
|
image_pipeline = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16) |
|
image_pipeline.enable_model_cpu_offload() |
|
|
|
@app.on_event("startup") |
|
async def startup_event(): |
|
global tokenizer, unified_model |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
tokenizer_data_bytes = redis_client.get(f"tokenizer:{tokenizer_name}") |
|
if tokenizer_data_bytes: |
|
tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
tokenizer.add_tokens(json.loads(tokenizer_data_bytes.decode("utf-8"))) |
|
tokenizer.pad_token = tokenizer.eos_token |
|
else: |
|
tokenizer = AutoTokenizer.from_pretrained("gpt2") |
|
tokenizer.pad_token = tokenizer.eos_token |
|
unified_model = UnifiedModel.load_model_from_redis(redis_client) |
|
unified_model.to(torch.device("cpu")) |
|
|
|
@app.post("/process") |
|
async def process(request: Request): |
|
global tokenizer, unified_model |
|
data = await request.json() |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
|
|
if data.get("train"): |
|
user_data = data.get("user_data", []) |
|
if not user_data: |
|
user_data = [ |
|
{"text": "Hola", "label": 1}, |
|
{"text": "Necesito ayuda", "label": 2}, |
|
{"text": "No entiendo", "label": 0} |
|
] |
|
redis_client.rpush("training_queue", json.dumps({ |
|
"tokenizers": {tokenizer_name: tokenizer.get_vocab()}, |
|
"data": user_data |
|
})) |
|
return {"message": "Training data received. Model will be updated asynchronously."} |
|
elif data.get("message"): |
|
user_id = data.get("user_id") |
|
text = data['message'] |
|
language = data.get("language", default_language) |
|
if user_id not in conversation_history: |
|
conversation_history[user_id] = [] |
|
conversation_history[user_id].append(text) |
|
contextualized_text = " ".join(conversation_history[user_id][-3:]) |
|
tokenized_input = tokenizer(contextualized_text, return_tensors="pt") |
|
with torch.no_grad(): |
|
logits = unified_model(**tokenized_input).logits |
|
predicted_class = torch.argmax(logits, dim=-1).item() |
|
response = chatbot_service.get_response(user_id, contextualized_text, language) |
|
redis_client.rpush("training_queue", json.dumps({ |
|
"tokenizers": {tokenizer_name: tokenizer.get_vocab()}, |
|
"data": [{"text": contextualized_text, "label": predicted_class}] |
|
})) |
|
return {"answer": response} |
|
else: |
|
raise HTTPException(status_code=400, detail="Request must contain 'train' or 'message'.") |
|
|
|
@app.get("/") |
|
async def get_home(): |
|
user_id = str(uuid.uuid4()) |
|
html_code = f""" |
|
<!DOCTYPE html> |
|
<html> |
|
<head> |
|
<meta charset="UTF-8"> |
|
<title>Chatbot</title> |
|
<style> |
|
body {{ |
|
font-family: 'Arial', sans-serif; |
|
background-color: #f4f4f9; |
|
margin: 0; |
|
padding: 0; |
|
display: flex; |
|
align-items: center; |
|
justify-content: center; |
|
min-height: 100vh; |
|
}} |
|
.container {{ |
|
background-color: #fff; |
|
border-radius: 10px; |
|
box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1); |
|
overflow: hidden; |
|
width: 400px; |
|
max-width: 90%; |
|
}} |
|
h1 {{ |
|
color: #333; |
|
text-align: center; |
|
padding: 20px; |
|
margin: 0; |
|
background-color: #f8f9fa; |
|
border-bottom: 1px solid #eee; |
|
}} |
|
#chatbox {{ |
|
height: 300px; |
|
overflow-y: auto; |
|
padding: 10px; |
|
border-bottom: 1px solid #eee; |
|
}} |
|
.message {{ |
|
margin-bottom: 10px; |
|
padding: 10px; |
|
border-radius: 5px; |
|
}} |
|
.message.user {{ |
|
background-color: #e1f5fe; |
|
text-align: right; |
|
}} |
|
.message.bot {{ |
|
background-color: #f1f1f1; |
|
text-align: left; |
|
}} |
|
#input {{ |
|
display: flex; |
|
padding: 10px; |
|
}} |
|
#input textarea {{ |
|
flex: 1; |
|
padding: 10px; |
|
border: 1px solid #ddd; |
|
border-radius: 4px; |
|
margin-right: 10px; |
|
}} |
|
#input button {{ |
|
padding: 10px 20px; |
|
border: none; |
|
border-radius: 4px; |
|
background-color: #007bff; |
|
color: #fff; |
|
cursor: pointer; |
|
}} |
|
#input button:hover {{ |
|
background-color: #0056b3; |
|
}} |
|
</style> |
|
</head> |
|
<body> |
|
<div class="container"> |
|
<h1>Chatbot</h1> |
|
<div id="chatbox"></div> |
|
<div id="input"> |
|
<textarea id="message" rows="3" placeholder="Escribe tu mensaje aquí..."></textarea> |
|
<button id="send">Enviar</button> |
|
</div> |
|
</div> |
|
<script> |
|
const chatbox = document.getElementById('chatbox'); |
|
const messageInput = document.getElementById('message'); |
|
const sendButton = document.getElementById('send'); |
|
|
|
function appendMessage(text, sender) {{ |
|
const messageDiv = document.createElement('div'); |
|
messageDiv.classList.add('message', sender); |
|
messageDiv.textContent = text; |
|
chatbox.appendChild(messageDiv); |
|
chatbox.scrollTop = chatbox.scrollHeight; |
|
}} |
|
|
|
async function sendMessage() {{ |
|
const message = messageInput.value; |
|
if (!message.trim()) return; |
|
|
|
appendMessage(message, 'user'); |
|
messageInput.value = ''; |
|
|
|
const response = await fetch('/process', {{ |
|
method: 'POST', |
|
headers: {{ |
|
'Content-Type': 'application/json' |
|
}}, |
|
body: JSON.stringify({{ |
|
message: message, |
|
user_id: '{user_id}' |
|
}}) |
|
}}); |
|
const data = await response.json(); |
|
appendMessage(data.answer, 'bot'); |
|
}} |
|
|
|
sendButton.addEventListener('click', sendMessage); |
|
messageInput.addEventListener('keypress', (e) => {{ |
|
if (e.key === 'Enter' && !e.shiftKey) {{ |
|
e.preventDefault(); |
|
sendMessage(); |
|
}} |
|
}}); |
|
</script> |
|
</body> |
|
</html> |
|
""" |
|
return HTMLResponse(content=html_code) |
|
|
|
def train_unified_model(): |
|
global tokenizer, unified_model |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
while True: |
|
training_data = redis_client.lpop("training_queue") |
|
if training_data: |
|
item_data = json.loads(training_data) |
|
tokenizer_data = item_data["tokenizers"] |
|
tokenizer_name = list(tokenizer_data.keys())[0] |
|
if redis_client.exists(f"tokenizer:{tokenizer_name}"): |
|
tokenizer.add_tokens(list(tokenizer_data[tokenizer_name].keys())) |
|
data = item_data["data"] |
|
dataset = SyntheticDataset(tokenizer, data) |
|
|
|
model_name = "unified_model" |
|
model_path = f"models/{model_name}" |
|
|
|
training_args = TrainingArguments( |
|
output_dir="./results", |
|
per_device_train_batch_size=8, |
|
num_train_epochs=3, |
|
) |
|
trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset) |
|
trainer.train() |
|
unified_model.save_pretrained(model_path) |
|
|
|
async def auto_learn(): |
|
global tokenizer, unified_model |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
while True: |
|
training_data = redis_client.lpop("training_queue") |
|
if training_data: |
|
item_data = json.loads(training_data) |
|
tokenizer_data = item_data["tokenizers"] |
|
tokenizer_name = list(tokenizer_data.keys())[0] |
|
if redis_client.exists(f"tokenizer:{tokenizer_name}"): |
|
tokenizer.add_tokens(list(tokenizer_data[tokenizer_name].keys())) |
|
data = item_data["data"] |
|
dataset = SyntheticDataset(tokenizer, data) |
|
|
|
model_name = "unified_model" |
|
model_path = f"models/{model_name}" |
|
|
|
training_args = TrainingArguments( |
|
output_dir="./results", |
|
per_device_train_batch_size=8, |
|
num_train_epochs=3, |
|
) |
|
trainer = Trainer(model=unified_model, args=training_args, train_dataset=dataset) |
|
trainer.train() |
|
unified_model.save_pretrained(model_path) |
|
|
|
async def auto_learn_music(): |
|
global musicgen_tokenizer, musicgen_model |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
while True: |
|
music_training_data = redis_client.lpop("music_training_queue") |
|
if music_training_data: |
|
music_training_data = json.loads(music_training_data.decode("utf-8")) |
|
inputs = musicgen_tokenizer(music_training_data, return_tensors="pt", padding=True) |
|
musicgen_model.train() |
|
optimizer = torch.optim.Adam(musicgen_model.parameters(), lr=5e-5) |
|
loss_fn = torch.nn.CrossEntropyLoss() |
|
|
|
for epoch in range(1): |
|
outputs = musicgen_model(**inputs) |
|
loss = loss_fn(outputs.logits, inputs['labels']) |
|
optimizer.zero_grad() |
|
loss.backward() |
|
optimizer.step() |
|
|
|
async def auto_learn_images(): |
|
global image_pipeline |
|
redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD) |
|
while True: |
|
image_training_data = redis_client.lpop("image_training_queue") |
|
if image_training_data: |
|
image_training_data = json.loads(image_training_data.decode("utf-8")) |
|
for image_prompt in image_training_data: |
|
image = image_pipeline( |
|
image_prompt, |
|
guidance_scale=0.0, |
|
num_inference_steps=4, |
|
max_sequence_length=256, |
|
generator=torch.Generator("cpu").manual_seed(0) |
|
).images[0] |
|
image_tensor = torch.tensor(np.array(image)).unsqueeze(0) |
|
image_pipeline.model.train() |
|
optimizer = torch.optim.Adam(image_pipeline.model.parameters(), lr=1e-5) |
|
loss_fn = torch.nn.MSELoss() |
|
target_tensor = torch.zeros_like(image_tensor) |
|
for epoch in range(1): |
|
outputs = image_pipeline.model(image_tensor) |
|
loss = loss_fn(outputs, target_tensor) |
|
optimizer.zero_grad() |
|
loss.backward() |
|
optimizer.step() |
|
|
|
|
|
if __name__ == "__main__": |
|
training_process = multiprocessing.Process(target=train_unified_model) |
|
training_process.start() |
|
music_training_process = multiprocessing.Process(target=auto_learn_music) |
|
music_training_process.start() |
|
image_training_process = multiprocessing.Process(target=auto_learn_images) |
|
image_training_process.start() |
|
import uvicorn |
|
uvicorn.run(app, host="0.0.0.0", port=7860) |
