word_train.py

#!/usr/bin/env python3
"""
Working Neural Chatbot - Word-Level seq2seq
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import pickle
import os

DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
BATCH_SIZE = 4
LEARNING_RATE = 0.001
EPOCHS = 200
MAX_LEN = 20
EMBED_DIM = 64
HIDDEN_DIM = 128

class WordVocab:
    """Word-level tokenization"""
    def __init__(self):
        self.word2idx = {'<PAD>': 0, '<UNK>': 1}
        self.idx2word = {0: '<PAD>', 1: '<UNK>'}
        self.next_idx = 2
    
    def add_word(self, word):
        if word not in self.word2idx:
            self.word2idx[word] = self.next_idx
            self.idx2word[self.next_idx] = word
            self.next_idx += 1
    
    def encode(self, text):
        words = text.lower().split()
        return [self.word2idx.get(w, 1) for w in words[:MAX_LEN]]
    
    def decode(self, tokens):
        words = [self.idx2word.get(t, '<UNK>') for t in tokens if t not in [0, 1]]
        return ' '.join(words)
    
    def __len__(self):
        return self.next_idx

class ChatDataset(Dataset):
    def __init__(self, conversations, vocab):
        self.pairs = []
        for q, a in conversations:
            q_tokens = vocab.encode(q)
            a_tokens = vocab.encode(a)
            q_padded = q_tokens + [0] * (MAX_LEN - len(q_tokens))
            a_padded = a_tokens + [0] * (MAX_LEN - len(a_tokens))
            self.pairs.append((q_padded[:MAX_LEN], a_padded[:MAX_LEN]))
    
    def __len__(self):
        return len(self.pairs)
    
    def __getitem__(self, idx):
        q, a = self.pairs[idx]
        return torch.tensor(q, dtype=torch.long), torch.tensor(a, dtype=torch.long)

class Seq2Seq(nn.Module):
    def __init__(self, vocab_size, embed_dim, hidden_dim):
        super().__init__()
        self.vocab_size = vocab_size
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.encoder = nn.LSTM(embed_dim, hidden_dim, num_layers=2, batch_first=True)
        self.decoder = nn.LSTM(embed_dim, hidden_dim, num_layers=2, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
    
    def forward(self, questions, answers):
        q_embed = self.embed(questions)
        _, (hidden, cell) = self.encoder(q_embed)
        a_embed = self.embed(answers)
        decoder_out, _ = self.decoder(a_embed, (hidden, cell))
        logits = self.fc(decoder_out)
        return logits
    
    def generate(self, question_tokens, max_len=10, device='cpu'):
        """Generate response token by token - FIXED"""
        self.eval()
        with torch.no_grad():
            # Encode question
            q_input = torch.tensor([question_tokens], dtype=torch.long).to(device)
            q_embed = self.embed(q_input)
            _, (hidden, cell) = self.encoder(q_embed)
            
            # Decode - start with PAD token
            answer_tokens = []
            current_token = torch.tensor([[0]], dtype=torch.long).to(device)
            
            h, c = hidden, cell
            
            for step in range(max_len):
                # Embed current token
                token_embed = self.embed(current_token)
                
                # Decoder step - properly threads hidden state
                decoder_out, (h, c) = self.decoder(token_embed, (h, c))
                
                # Get logits for next token
                logits = self.fc(decoder_out[0, 0, :])
                
                # Greedy for first few tokens, then sample
                if step < 2:
                    next_token = torch.argmax(logits).item()
                else:
                    probs = torch.softmax(logits / 0.7, dim=0)
                    next_token = torch.multinomial(probs, 1).item()
                
                # Stop if PAD or UNK token
                if next_token in [0, 1]:
                    break
                
                answer_tokens.append(next_token)
                current_token = torch.tensor([[next_token]], dtype=torch.long).to(device)
            
            return answer_tokens

def get_data():
    return [
        ("hello", "hi"),
        ("hi", "hello"),
        ("how are you", "great"),
        ("what is your name", "chatbot"),
        ("goodbye", "bye"),
        ("thanks", "welcome"),
        ("help", "help you"),
        ("good morning", "morning"),
        ("how is it going", "good"),
        ("tell me a joke", "joke"),
        ("are you smart", "yes"),
        ("who are you", "ai"),
        ("what can you do", "chat"),
        ("bye", "later"),
        ("cool", "thanks"),
        ("like you", "like you too"),
    ]

def train():
    print("="*60)
    print("Training Word-Level Neural Chatbot")
    print("="*60)
    
    # Build vocab
    print("\n1. Building vocabulary...")
    convs = get_data()
    vocab = WordVocab()
    for q, a in convs:
        for word in q.lower().split() + a.lower().split():
            vocab.add_word(word)
    print(f"   Vocab size: {len(vocab)}")
    
    # Dataset
    print("\n2. Preparing dataset...")
    dataset = ChatDataset(convs, vocab)
    loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
    print(f"   Samples: {len(dataset)}")
    
    # Model
    print("\n3. Creating model...")
    model = Seq2Seq(len(vocab), EMBED_DIM, HIDDEN_DIM).to(DEVICE)
    opt = optim.Adam(model.parameters(), lr=LEARNING_RATE)
    loss_fn = nn.CrossEntropyLoss(ignore_index=0)
    
    # Train
    print("\n4. Training...")
    model.train()
    for epoch in range(EPOCHS):
        total_loss = 0
        for q, a in loader:
            q, a = q.to(DEVICE), a.to(DEVICE)
            logits = model(q, a)
            loss = loss_fn(logits.view(-1, len(vocab)), a.view(-1))
            opt.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            total_loss += loss.item()
        
        if (epoch + 1) % 20 == 0:
            print(f"   Epoch {epoch+1}/{EPOCHS}, Loss: {total_loss/len(loader):.4f}")
    
    # Save
    print("\n5. Saving...")
    torch.save({
        'model': model.state_dict(),
        'config': {'vocab_size': len(vocab), 'embed_dim': EMBED_DIM, 'hidden_dim': HIDDEN_DIM}
    }, 'word_model.pt')
    
    with open('word_vocab.pkl', 'wb') as f:
        pickle.dump(vocab, f)
    
    print("   Saved!")
    print("="*60)

if __name__ == "__main__":
    train()