Create model.py

model.py

@@ -0,0 +1,169 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from tqdm import tqdm
+import math
+
+# 1. Dataset class for loading and processing data
+class FullChatDataset(Dataset):
+    def __init__(self, dataset_names=["blended_skill_talk", "conv_ai_2", "social_i_qa"], max_length=128):
+        self.datasets = []
+        
+        # Load all specified datasets
+        for name in dataset_names:
+            try:
+                dataset = load_dataset(name, split="train")
+                self.datasets.append(dataset)
+            except Exception as e:
+                print(f"Failed to load dataset {name}: {e}")
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+        self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        self.max_length = max_length
+        
+    def __len__(self):
+        return sum(len(d) for d in self.datasets)
+    
+    def __getitem__(self, idx):
+        # Determine which dataset the index belongs to
+        for dataset in self.datasets:
+            if idx < len(dataset):
+                item = dataset[idx]
+                break
+            idx -= len(dataset)
+        
+        # Handling different dataset formats
+        if 'dialog' in item:  # For Daily Dialog
+            dialog = item['dialog']
+        elif 'messages' in item:  # For some other datasets
+            dialog = [msg['text'] for msg in item['messages']]
+        else:  # Universal handling
+            dialog = [v for k, v in item.items() if isinstance(v, str)]
+        
+        context = " [SEP] ".join(dialog[:-1])
+        response = dialog[-1]
+        
+        inputs = self.tokenizer(
+            context,
+            text_pair=response,
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        return {
+            'input_ids': inputs['input_ids'].flatten(),
+            'attention_mask': inputs['attention_mask'].flatten(),
+            'labels': inputs['input_ids'].flatten()
+        }
+
+# 2. Model architecture
+class SimpleTransformerModel(nn.Module):
+    def __init__(self, vocab_size, d_model=256, nhead=4, num_layers=3):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model, nhead)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers)
+        self.fc = nn.Linear(d_model, vocab_size)
+        
+    def forward(self, x, mask=None):
+        x = self.embedding(x)
+        x = self.pos_encoder(x)
+        x = self.transformer(x, mask)
+        return self.fc(x)
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=500):
+        super().__init__()
+        position = torch.arange(max_len).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(max_len, d_model)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        return x + self.pe[:x.size(1)]
+
+# 3. Model training
+def train(model, dataloader, epochs=3, lr=3e-4):
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+    criterion = nn.CrossEntropyLoss(ignore_index=0)
+    optimizer = optim.Adam(model.parameters(), lr=lr)
+    
+    for epoch in range(epochs):
+        model.train()
+        total_loss = 0
+        pbar = tqdm(dataloader, desc=f"Epoch {epoch+1}/{epochs}")
+        
+        for batch in pbar:
+            inputs = batch['input_ids'].to(device)
+            masks = batch['attention_mask'].to(device)
+            labels = batch['labels'].to(device)
+            
+            optimizer.zero_grad()
+            outputs = model(inputs, masks)
+            loss = criterion(outputs.view(-1, outputs.size(-1)), labels.view(-1))
+            loss.backward()
+            optimizer.step()
+            
+            total_loss += loss.item()
+            pbar.set_postfix({'loss': loss.item()})
+        
+        print(f"Epoch {epoch+1} - Avg loss: {total_loss/len(dataloader):.4f}")
+
+# 4. Response generation
+def chat(model, tokenizer, prompt, max_length=50):
+    device = next(model.parameters()).device
+    model.eval()
+    
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        max_length=128,
+        truncation=True,
+        padding='max_length'
+    ).to(device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            input_ids=inputs['input_ids'],
+            attention_mask=inputs['attention_mask'],
+            max_length=max_length,
+            do_sample=True,
+            top_k=50,
+            top_p=0.95,
+            temperature=0.7
+        )
+    
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# 5. Main process
+if __name__ == "__main__":
+    # Initialization
+    dataset = FullChatDataset()
+    dataloader = DataLoader(dataset, batch_size=16, shuffle=True)
+    
+    # Model creation
+    model = SimpleTransformerModel(len(dataset.tokenizer))
+    
+    # Training
+    train(model, dataloader)
+    
+    # Saving
+    torch.save(model.state_dict(), "chatbot_model.pt")
+    dataset.tokenizer.save_pretrained("chatbot_tokenizer")
+    
+    
+    while True:
+        user_input = input("You: ")
+        if user_input.lower() in ['exit', 'quit']:
+            break
+        response = chat(model, dataset.tokenizer, user_input)
+        print(f"Bot: {response}")