Update main.py

main.py

@@ -46,6 +46,17 @@ class UnifiedModel(nn.Module):
         logits = self.classifier(concatenated_hidden_states)
         return logits
 
+    @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=2)
+            model.load_state_dict(torch.load(model_data_bytes))
+        else:
+            model = AutoModelForSequenceClassification.from_pretrained("gpt2")
+        return UnifiedModel([model])
+
 class SyntheticDataset(Dataset):
     def __init__(self, tokenizers, data):
         self.tokenizers = tokenizers
@@ -68,6 +79,20 @@ class SyntheticDataset(Dataset):
 
 @app.post("/process")
 async def process(request: Request):
+    """
+    Processes requests for training and prediction.
+
+    Args:
+        request (Request): The incoming request object.
+
+    Returns:
+        dict: A dictionary containing either a message indicating successful
+              training data submission or the model's prediction.
+
+    Raises:
+        HTTPException: If the request does not contain 'train' or 'predict'
+                       keys.
+    """
     data = await request.json()
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     
@@ -102,41 +127,13 @@ async def process(request: Request):
         if not user_data:
             user_data = [{"text": "Sample text for automatic training.", "label": 0}]
         
-        train_dataset = SyntheticDataset(tokenizers, user_data)
-        train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-        
-        training_args = TrainingArguments(
-            output_dir="memory",
-            evaluation_strategy="epoch",
-            learning_rate=5e-5,
-            per_device_train_batch_size=8,
-            per_device_eval_batch_size=8,
-            num_train_epochs=10,
-            weight_decay=0.01,
-            logging_steps=10,
-            optim="adamw_hf"
-        )
-        
-        optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
-        unified_model.train()
-        
-        for epoch in range(training_args.num_train_epochs):
-            for batch in train_loader:
-                input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
-                attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
-                labels = batch["labels"].to("cpu")
-                outputs = unified_model(input_ids=input_ids, attention_mask=attention_mask)
-                loss = nn.CrossEntropyLoss()(outputs, labels)
-                loss.backward()
-                optimizer.step()
-                optimizer.zero_grad()
-                
-                print(f"Epoch {epoch}, Loss {loss.item()}")
-        
-        print("Training complete.")
+        # Add user data to Redis queue for asynchronous training
+        redis_client.rpush("training_queue", json.dumps({
+            "tokenizers": {tokenizer_name: tokenizer.get_vocab()},
+            "data": user_data
+        }))
 
-        push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name)
-        return {"message": "Model trained and updated in Redis."}
+        return {"message": "Training data received. Model will be updated asynchronously."}
     
     elif data.get("predict"):
         text = data['text']
@@ -155,6 +152,19 @@ async def process(request: Request):
 
 @app.post("/external_answer")
 async def external_answer(request: Request):
+    """
+    Provides an answer to a question using the unified model and triggers
+    asynchronous training with the new question-answer pair.
+
+    Args:
+        request (Request): The incoming request object containing the question.
+
+    Returns:
+        dict: A dictionary containing the answer to the question.
+
+    Raises:
+        HTTPException: If the request does not contain a 'question' key.
+    """
     data = await request.json()
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     
@@ -162,26 +172,16 @@ async def external_answer(request: Request):
     if not question:
         raise HTTPException(status_code=400, detail="Question is required.")
 
-    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=2)
-        model.load_state_dict(torch.load(model_data_bytes))
-    else:
-        model = AutoModelForSequenceClassification.from_pretrained("gpt2")
+    # Load the model and tokenizer from Redis
+    unified_model = UnifiedModel.load_model_from_redis(redis_client)
+    unified_model.to(torch.device("cpu"))
 
+    tokenizer_data_bytes = redis_client.get(f"tokenizer:unified_tokenizer")
     if tokenizer_data_bytes:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         tokenizer.add_tokens(json.loads(tokenizer_data_bytes))
     else:
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
-
-    unified_model = UnifiedModel([model])
-    unified_model.to(torch.device("cpu"))
     
     tokenized_input = tokenizer(question, return_tensors="pt")
     input_ids = tokenized_input['input_ids']
@@ -192,45 +192,22 @@ async def external_answer(request: Request):
         predicted_class = torch.argmax(logits, dim=-1).item()
     response = {"answer": f"Response to '{question}' is class {predicted_class}"}
     
-    extreme_training_data = [{"text": question, "label": predicted_class}]
-    train_dataset = SyntheticDataset({tokenizer_name: tokenizer}, extreme_training_data)
-    train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
-    
-    training_args = TrainingArguments(
-        output_dir="memory",
-        evaluation_strategy="epoch",
-        learning_rate=5e-5,
-        per_device_train_batch_size=8,
-        per_device_eval_batch_size=8,
-        num_train_epochs=10,
-        weight_decay=0.01,
-        logging_steps=10,
-        optim="adamw_hf"
-    )
-    
-    optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
-    unified_model.train()
-    
-    for epoch in range(training_args.num_train_epochs):
-        for batch in train_loader:
-            input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in [tokenizer_name]]
-            attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in [tokenizer_name]]
-            labels = batch["labels"].to("cpu")
-            outputs = unified_model(input_ids=input_ids, attention_mask=attention_mask)
-            loss = nn.CrossEntropyLoss()(outputs, labels)
-            loss.backward()
-            optimizer.step()
-            optimizer.zero_grad()
-            
-            print(f"Epoch {epoch}, Loss {loss.item()}")
-    
-    print("Extreme training complete.")
-    push_to_redis({model_name: model}, {tokenizer_name: tokenizer}, redis_client, model_name, tokenizer_name)
+    # Asynchronously train on the new data point
+    redis_client.rpush("training_queue", json.dumps({
+        "tokenizers": {"unified_tokenizer": tokenizer.get_vocab()},
+        "data": [{"text": question, "label": predicted_class}]
+    }))
     
     return response
 
 @app.get("/")
 async def get_home():
+    """
+    Serves a basic HTML page as the home route.
+
+    Returns:
+        HTMLResponse: The HTML content of the home page.
+    """
     html_code = """
     <!DOCTYPE html>
     <html>
@@ -264,6 +241,18 @@ async def get_home():
     return HTMLResponse(content=html_code)
 
 def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
+    """
+    Saves the given models and tokenizers to Redis.
+
+    Args:
+        models (dict): A dictionary of model names and their corresponding
+                       PyTorch models.
+        tokenizers (dict): A dictionary of tokenizer names and their
+                           corresponding tokenizers.
+        redis_client: The Redis client instance.
+        model_name (str): The base name to use for saving the models.
+        tokenizer_name (str): The base name to use for saving the tokenizers.
+    """
     for model_name, model in models.items():
         torch.save(model.state_dict(), model_name)
         redis_client.set(f"model:{model_name}", open(model_name, "rb").read())
@@ -273,11 +262,14 @@ def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
         redis_client.set(f"tokenizer:{tokenizer_name}", json.dumps(tokens))
 
 def continuous_training():
+    """
+    Continuously checks for new training data in Redis and updates the model.
+    """
     redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
     
     while True:
         try:
-            data = redis_client.get("training_queue")
+            data = redis_client.lpop("training_queue")
             if data:
                 data = json.loads(data)
                 unified_model = UnifiedModel.load_model_from_redis(redis_client)
@@ -302,13 +294,16 @@ def continuous_training():
                         print(f"Epoch {epoch}, Loss {loss.item()}")
                 
                 push_to_redis(unified_model.models, data["tokenizers"], redis_client, "unified_model", "unified_tokenizer")
-                redis_client.delete("training_queue")
             time.sleep(10)
         except Exception as e:
             print(f"Error in continuous training: {e}")
             time.sleep(5)
 
 if __name__ == "__main__":
-    continuous_training()
+    # Start the continuous training process in a separate process
+    training_process = multiprocessing.Process(target=continuous_training)
+    training_process.start()
+
+    # Run the FastAPI app
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)