Create App.py

App.py

@@ -0,0 +1,428 @@
+from dotenv import load_dotenv
+import os
+import json
+import requests
+import redis
+from transformers import (
+    AutoTokenizer,
+    AutoModel,
+    TrainingArguments,
+)
+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 pydantic import BaseModel
+from typing import List, Dict
+from fastapi.responses import HTMLResponse
+import multiprocessing
+import time
+
+load_dotenv()
+
+REDIS_HOST = os.getenv('REDIS_HOST')
+REDIS_PORT = os.getenv('REDIS_PORT')
+REDIS_PASSWORD = os.getenv('REDIS_PASSWORD')
+
+app = FastAPI()
+
+class UnifiedModel(nn.Module):
+    def __init__(self, models):
+        super(UnifiedModel, self).__init__()
+        self.models = nn.ModuleList(models)
+        self.classifier = nn.Linear(sum([model.config.hidden_size for model in models]), 2)
+
+    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.last_hidden_state[:, 0, :])
+        concatenated_hidden_states = torch.cat(hidden_states, dim=-1)
+        logits = self.classifier(concatenated_hidden_states)
+        return logits
+
+class SyntheticDataset(Dataset):
+    def __init__(self, tokenizers, data):
+        self.tokenizers = tokenizers
+        self.data = data
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        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["label"] = torch.tensor(label)
+        return tokenized
+
+@app.post("/process")
+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_data = json.loads(model_data_bytes)
+        model = AutoModel.from_pretrained("gpt2")
+        model.load_state_dict(torch.load(model_data))
+    else:
+        model = AutoModel.from_pretrained("gpt2")
+    models[model_name] = model
+
+    if tokenizer_data_bytes:
+        tokenizer_data = json.loads(tokenizer_data_bytes)
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.add_tokens(tokenizer_data)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+    tokenizers[tokenizer_name] = tokenizer
+
+    unified_model = UnifiedModel(list(models.values()))
+    unified_model.to(torch.device("cpu"))
+    
+    if data.get("train"):
+        user_data = data.get("user_data", [])
+        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["label"].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.")
+
+        push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name)
+        return {"message": "Model trained and updated in Redis."}
+    
+    elif data.get("predict"):
+        text = data['text']
+        tokenized_inputs = [tokenizers[name](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]
+        
+        with torch.no_grad():
+            logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
+            predicted_class = torch.argmax(logits, dim=-1).item()
+        
+        return {"prediction": predicted_class}
+    
+    else:
+        raise HTTPException(status_code=400, detail="Request must contain 'train' or 'predict'.")
+
+@app.post("/external_answer")
+async def external_answer(request: Request):
+    data = await request.json()
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
+    
+    question = data.get('question')
+    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_data = json.loads(model_data_bytes)
+        model = AutoModel.from_pretrained("gpt2")
+        model.load_state_dict(torch.load(model_data))
+    else:
+        model = AutoModel.from_pretrained("gpt2")
+
+    if tokenizer_data_bytes:
+        tokenizer_data = json.loads(tokenizer_data_bytes)
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        tokenizer.add_tokens(tokenizer_data)
+    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']
+    attention_mask = tokenized_input['attention_mask']
+    
+    with torch.no_grad():
+        logits = unified_model(input_ids=input_ids, attention_mask=attention_mask)
+        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["label"].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)
+    
+    return response
+
+@app.get("/")
+async def get_home():
+    html_code = """
+    <!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;
+            }
+            .container {
+                max-width: 1200px;
+                margin: 0 auto;
+                padding: 20px;
+            }
+            h1 {
+                color: #333;
+                text-align: center;
+            }
+            .grid-container {
+                display: grid;
+                grid-template-columns: repeat(auto-fill, minmax(300px, 1fr));
+                gap: 10px;
+                margin-top: 20px;
+            }
+            .grid-item {
+                background: #fff;
+                padding: 20px;
+                border-radius: 8px;
+                box-shadow: 0 0 10px rgba(0, 0, 0, 0.1);
+                transition: transform 0.3s;
+            }
+            .grid-item:hover {
+                transform: scale(1.05);
+            }
+            .question {
+                font-weight: bold;
+                color: #007bff;
+            }
+            .answer {
+                margin-top: 10px;
+                color: #333;
+            }
+            input[type="text"] {
+                width: calc(100% - 22px);
+                padding: 10px;
+                margin: 0;
+                border: 1px solid #ddd;
+                border-radius: 4px;
+            }
+            button {
+                padding: 10px 20px;
+                background-color: #007bff;
+                color: #fff;
+                border: none;
+                border-radius: 4px;
+                cursor: pointer;
+                margin-top: 10px;
+            }
+            button:hover {
+                background-color: #0056b3;
+            }
+        </style>
+        <script>
+            async function sendMessage() {
+                const question = document.getElementById('question').value;
+                const responseElement = document.getElementById('response');
+                
+                const response = await fetch('/external_answer', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({ question: question })
+                });
+                
+                const data = await response.json();
+                responseElement.innerText = "Response: " + data.answer;
+                
+                const gridContainer = document.getElementById('grid-container');
+                const newItem = document.createElement('div');
+                newItem.classList.add('grid-item');
+                newItem.innerHTML = `<div class="question">${question}</div><div class="answer">${data.answer}</div>`;
+                gridContainer.prepend(newItem);
+            }
+        </script>
+    </head>
+    <body>
+        <div class="container">
+            <h1>Chatbot</h1>
+            <input type="text" id="question" placeholder="Ask me something...">
+            <button onclick="sendMessage()">Send</button>
+            <div id="response"></div>
+            <div class="grid-container" id="grid-container"></div>
+        </div>
+    </body>
+    </html>
+    """
+    return HTMLResponse(content=html_code)
+
+def push_to_redis(models, tokenizers, redis_client, model_name, tokenizer_name):
+    model_data = json.dumps(next(iter(models.values())).state_dict())
+    redis_client.set(f"model:{model_name}", model_data)
+
+    tokenizer_data = json.dumps(next(iter(tokenizers.values())).get_vocab())
+    redis_client.set(f"tokenizer:{tokenizer_name}", tokenizer_data)
+
+def continuous_training():
+    redis_client = redis.StrictRedis(host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True)
+
+    model_name = "unified_model"
+    tokenizer_name = "unified_tokenizer"
+
+    while True:
+        try:
+            model_data_bytes = redis_client.get(f"model:{model_name}")
+            tokenizer_data_bytes = redis_client.get(f"tokenizer:{tokenizer_name}")
+
+            if model_data_bytes and tokenizer_data_bytes:
+                model_data = json.loads(model_data_bytes)
+                model = AutoModel.from_pretrained("gpt2")
+                model.load_state_dict(torch.load(model_data))
+
+                tokenizer_data = json.loads(tokenizer_data_bytes)
+                tokenizer = AutoTokenizer.from_pretrained("gpt2")
+                tokenizer.add_tokens(tokenizer_data)
+
+                unified_model = UnifiedModel([model])
+                unified_model.to(torch.device("cpu"))
+
+                train_data = [{"text": "Sample training text.", "label": 0}]
+                train_dataset = SyntheticDataset({tokenizer_name: tokenizer}, train_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["label"].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.")
+                push_to_redis({model_name: model}, {tokenizer_name: tokenizer}, redis_client, model_name, tokenizer_name)
+            else:
+                print("No model or tokenizer found in Redis. Skipping training.")
+
+            time.sleep(600)
+
+        except Exception as e:
+            print(f"An error occurred: {e}")
+            time.sleep(60)
+
+def start_server():
+    import uvicorn
+    cpu_cores = os.cpu_count() or 1
+    num_workers = max(1, cpu_cores - 1)
+    
+    uvicorn.run(app, host="0.0.0.0", port=8000, workers=num_workers, timeout_keep_alive=0)
+
+if __name__ == "__main__":
+    api_process = multiprocessing.Process(target=start_server)
+    training_process = multiprocessing.Process(target=continuous_training)
+    
+    api_process.start()
+    training_process.start()
+    
+    api_process.join()
+    training_process.join()