Update main.py

main.py

@@ -6,11 +6,9 @@ from transformers import (
     AutoTokenizer,
     AutoModelForSequenceClassification,
     AutoModelForCausalLM,
+    TrainingArguments,
+    Trainer,
 )
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader, Dataset
-from torch.optim import AdamW
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import HTMLResponse
 import multiprocessing
@@ -65,38 +63,26 @@ class ChatbotService:
 
 chatbot_service = ChatbotService()
 
-class UnifiedModel(nn.Module):
-    def __init__(self, models):
-        super(UnifiedModel, self).__init__()
-        self.models = nn.ModuleList(models)
-        hidden_size = self.models[0].config.hidden_size
-        self.projection = nn.Linear(len(models) * 3, 768)
-        self.classifier = nn.Linear(hidden_size, 3)
-
-    def forward(self, input_ids, attention_mask):
-        hidden_states = []
-        for model, input_id, attn_mask in zip(self.models, input_ids, attention_mask):
-            outputs = model(input_ids=input_id, attention_mask=attn_mask)
-            hidden_states.append(outputs.logits)
-        concatenated_hidden_states = torch.cat(hidden_states, dim=1)
-        projected_features = self.projection(concatenated_hidden_states)
-        logits = self.classifier(projected_features)
-        return logits
+class UnifiedModel(AutoModelForSequenceClassification):
+    def __init__(self, config):
+        super().__init__(config)
 
     @staticmethod
     def load_model_from_redis(redis_client):
         model_name = "unified_model"
-        model_data_bytes = redis_client.get(f"model:{model_name}")
-        if model_data_bytes:
-            model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
-            model.load_state_dict(torch.load(model_data_bytes))
+        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 = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
-        return UnifiedModel([model, model])
+            model = UnifiedModel.from_pretrained(model_path)
+        return model
 
 class SyntheticDataset(Dataset):
-    def __init__(self, tokenizers, data):
-        self.tokenizers = tokenizers
+    def __init__(self, tokenizer, data):
+        self.tokenizer = tokenizer
         self.data = data
 
     def __len__(self):
@@ -106,13 +92,8 @@ class SyntheticDataset(Dataset):
         item = self.data[idx]
         text = item['text']
         label = item['label']
-        tokenized = {}
-        for name, tokenizer in self.tokenizers.items():
-            tokens = tokenizer(text, padding="max_length", truncation=True, max_length=128)
-            tokenized[f"input_ids_{name}"] = torch.tensor(tokens["input_ids"])
-            tokenized[f"attention_mask_{name}"] = torch.tensor(tokens["attention_mask"])
-        tokenized["labels"] = torch.tensor(label)
-        return tokenized
+        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 = {}
 
@@ -121,22 +102,10 @@ async def process(request: Request):
     data = await request.json()
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
 
-    tokenizers = {}
-    models = {}
-
-    model_name = "unified_model"
     tokenizer_name = "unified_tokenizer"
 
-    model_data_bytes = redis_client.get(f"model:{model_name}")
     tokenizer_data_bytes = redis_client.get(f"tokenizer:{tokenizer_name}")
 
-    if model_data_bytes:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
-        model.load_state_dict(torch.load(model_data_bytes))
-    else:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)
-    models[model_name] = model
-
     if tokenizer_data_bytes:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
@@ -144,9 +113,8 @@ async def process(request: Request):
     else:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         tokenizer.pad_token = tokenizer.eos_token
-    tokenizers[tokenizer_name] = tokenizer
 
-    unified_model = UnifiedModel(list(models.values()))
+    unified_model = UnifiedModel.load_model_from_redis(redis_client)
     unified_model.to(torch.device("cpu"))
 
     if data.get("train"):
@@ -170,11 +138,9 @@ async def process(request: Request):
             conversation_history[user_id] = []
         conversation_history[user_id].append(text)
         contextualized_text = " ".join(conversation_history[user_id][-3:])
-        tokenized_inputs = [tokenizers[name](contextualized_text, return_tensors="pt") for name in tokenizers.keys()]
-        input_ids = [tokens['input_ids'] for tokens in tokenized_inputs]
-        attention_mask = [tokens['attention_mask'] for tokens in tokenized_inputs]
+        tokenized_input = tokenizer(contextualized_text, return_tensors="pt")
         with torch.no_grad():
-            logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
+            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({
@@ -327,35 +293,26 @@ def train_unified_model():
         if training_queue:
             for item in training_queue:
                 item_data = json.loads(item)
-                tokenizers = {name: AutoTokenizer.from_pretrained("gpt2") for name in item_data["tokenizers"]}
-                for tokenizer in tokenizers.values():
-                    tokenizer.pad_token = tokenizer.eos_token
+                tokenizer_data = item_data["tokenizers"]
+                tokenizer_name = list(tokenizer_data.keys())[0]
+                tokenizer = AutoTokenizer.from_pretrained("gpt2")
+                tokenizer.add_tokens(json.loads(tokenizer_data[tokenizer_name]))
+                tokenizer.pad_token = tokenizer.eos_token
                 data = item_data["data"]
-                dataset = SyntheticDataset(tokenizers, data)
-                dataloader = DataLoader(dataset, batch_size=8, shuffle=True)
-
-                model = UnifiedModel([AutoModelForSequenceClassification.from_pretrained("gpt2", num_labels=3)])
-                optimizer = AdamW(model.parameters(), lr=1e-5)
-                criterion = nn.CrossEntropyLoss()
-
-                for epoch in range(3):
-                    model.train()
-                    for batch in dataloader:
-                        input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers]
-                        attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers]
-                        labels = batch["labels"].to("cpu")
-
-                        optimizer.zero_grad()
-                        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
-                        loss = criterion(outputs, labels)
-                        loss.backward()
-                        optimizer.step()
-
-                model_data_path = "model_data.pt"
-                torch.save(model.state_dict(), model_data_path)
-                with open(model_data_path, "rb") as f:
-                    model_data_bytes = f.read()
-                redis_client.set(f"model:unified_model", model_data_bytes)
+                dataset = SyntheticDataset(tokenizer, data)
+
+                model_name = "unified_model"
+                model_path = f"models/{model_name}"
+                model = UnifiedModel.from_pretrained(model_path)
+                
+                training_args = TrainingArguments(
+                    output_dir="./results",
+                    per_device_train_batch_size=8,
+                    num_train_epochs=3,
+                )
+                trainer = Trainer(model=model, args=training_args, train_dataset=dataset)
+                trainer.train()
+                model.save_pretrained(model_path)
                 redis_client.delete("training_queue")
         time.sleep(60)