import torch
from torch.utils.data import Dataset, DataLoader
from torch.nn.utils.rnn import pad_sequence
from collections import Counter
from itertools import chain

# Example data
source_sentences = ["hello how are you", "what is your name"]
target_sentences = ["hi i am fine", "i am a chatbot"]

# Tokenization and vocabulary
# Tokenize and detokenize
def tokenize(sentence):
    return sentence.split()

def numericalize(sentences, vocab):
    return [[vocab.get(word, vocab['<unk>']) for word in tokenize(sentence)] for sentence in sentences]

def detokenize(tokens, vocab):
    reverse_vocab = {idx: word for word, idx in vocab.items()}
    return ' '.join(reverse_vocab.get(idx, '<unk>') for idx in tokens)

def build_vocab(sentences):
    flat_list = list(chain(*sentences))
    counter = Counter(flat_list)
    vocab = {word: i+1 for i, (word, _) in enumerate(counter.items())}
    vocab['<pad>'] = 0
    vocab['<unk>'] = len(vocab) + 1
    return vocab

source_tokens = tokenize(source_sentences)
target_tokens = tokenize(target_sentences)

source_vocab = build_vocab(source_tokens)
target_vocab = build_vocab(target_tokens)

def numericalize(sentences, vocab):
    return [[vocab.get(word, vocab['<unk>']) for word in sentence] for sentence in sentences]

source_numerical = numericalize(source_tokens, source_vocab)
target_numerical = numericalize(target_tokens, target_vocab)

# Dataset
class Seq2SeqDataset(Dataset):
    def __init__(self, source_sentences, target_sentences):
        self.source_sentences = source_sentences
        self.target_sentences = target_sentences

    def __len__(self):
        return len(self.source_sentences)

    def __getitem__(self, idx):
        src = torch.tensor(self.source_sentences[idx], dtype=torch.long)
        trg = torch.tensor(self.target_sentences[idx], dtype=torch.long)
        return src, trg

def collate_fn(batch):
    src_batch, trg_batch = zip(*batch)
    src_lengths = [len(seq) for seq in src_batch]
    trg_lengths = [len(seq) for seq in trg_batch]

    src_padded = pad_sequence(src_batch, padding_value=source_vocab['<pad>'])
    trg_padded = pad_sequence(trg_batch, padding_value=target_vocab['<pad>'])

    return src_padded, trg_padded
import torch.nn as nn
import torch.nn.functional as F

class Encoder(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)

    def forward(self, x):
        embedded = self.embedding(x)
        outputs, (hidden, cell) = self.lstm(embedded)
        return hidden, cell

class Attention(nn.Module):
    def __init__(self, hidden_dim):
        super(Attention, self).__init__()
        self.Wa = nn.Linear(hidden_dim, hidden_dim)
        self.Ua = nn.Linear(hidden_dim, hidden_dim)
        self.Va = nn.Linear(hidden_dim, 1)

    def forward(self, query, keys, values):
        scores = self.Va(torch.tanh(self.Wa(query) + self.Ua(keys)))
        weights = F.softmax(scores, dim=-1)
        context = torch.bmm(weights, values)
        return context, weights

class Decoder(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.lstm = nn.LSTM(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.attention = Attention(hidden_dim)

    def forward(self, x, hidden, cell, encoder_outputs):
        embedded = self.embedding(x)
        outputs, (hidden, cell) = self.lstm(embedded, (hidden, cell))
        context, _ = self.attention(hidden[-1].unsqueeze(1), encoder_outputs, encoder_outputs)
        combined = torch.cat((outputs, context), dim=-1)
        predictions = self.fc(combined)
        return predictions, hidden, cell

class Seq2Seq(nn.Module):
    def __init__(self, encoder, decoder):
        super(Seq2Seq, self).__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, src, trg, teacher_forcing_ratio=0.5):
        batch_size = src.size(0)
        max_len = trg.size(1)
        vocab_size = self.decoder.fc.out_features

        outputs = torch.zeros(batch_size, max_len, vocab_size).to(src.device)
        encoder_hidden, encoder_cell = self.encoder(src)
        encoder_outputs, _ = self.encoder.lstm(self.encoder.embedding(src))

        input = trg[:, 0]
        for t in range(1, max_len):
            output, encoder_hidden, encoder_cell = self.decoder(input.unsqueeze(1), encoder_hidden, encoder_cell, encoder_outputs)
            outputs[:, t, :] = output.squeeze(1)
            top1 = output.argmax(2)
            input = trg[:, t] if torch.rand(1).item() < teacher_forcing_ratio else top1.squeeze(1)

        return outputs
import torch.optim as optim

def train(model, dataloader, criterion, optimizer, device):
    model.train()
    total_loss = 0
    for src, trg in dataloader:
        src, trg = src.to(device), trg.to(device)

        optimizer.zero_grad()
        output = model(src, trg)
        output_dim = output.shape[-1]
        trg = trg[:, 1:].contiguous().view(-1)
        output = output[:, 1:, :].contiguous().view(-1, output_dim)
        
        loss = criterion(output, trg)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
    return total_loss / len(dataloader)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
encoder = Encoder(vocab_size=len(source_vocab), embedding_dim=256, hidden_dim=512).to(device)
decoder = Decoder(vocab_size=len(target_vocab), embedding_dim=256, hidden_dim=512).to(device)
model = Seq2Seq(encoder, decoder).to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()

# Define the number of epochs
num_epochs = 10

# DataLoader
dataset = Seq2SeqDataset(source_numerical, target_numerical)
train_loader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)

# Training Loop
for epoch in range(num_epochs):
    train_loss = train(model, train_loader, criterion, optimizer, device)
    print(f'Epoch {epoch}, Loss: {train_loss:.4f}')
def generate_response(prompt, model, source_vocab, target_vocab, device, max_length=50):
    model.eval()
    with torch.no_grad():
        # Tokenize and numericalize prompt
        prompt_tokens = tokenize(prompt)
        prompt_numerical = [source_vocab.get(word, source_vocab['<unk>']) for word in prompt_tokens]
        src_tensor = torch.tensor(prompt_numerical).unsqueeze(0).to(device)

        # Generate response
        output = model(src_tensor, torch.zeros_like(src_tensor))
        output_tokens = output.argmax(-1).squeeze().cpu().numpy()

        # Detokenize output
        response = detokenize(output_tokens, target_vocab)
    return response

import signal

# Initialize vocabularies and model
source_vocab = {'<pad>': 0, '<unk>': 1}
target_vocab = {'<pad>': 0, '<unk>': 1}

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
encoder = Encoder(vocab_size=len(source_vocab), embedding_dim=256, hidden_dim=512).to(device)
decoder = Decoder(vocab_size=len(target_vocab), embedding_dim=256, hidden_dim=512).to(device)
model = Seq2Seq(encoder, decoder).to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
def handle_interrupt(signal, frame):
    print("Interrupt received, training model...")
    
    # Prepare data
    source_numerical = numericalize(source_sentences, source_vocab)
    target_numerical = numericalize(target_sentences, target_vocab)
    dataset = Seq2SeqDataset(source_numerical, target_numerical)
    dataloader = DataLoader(dataset, batch_size=2, shuffle=True, collate_fn=collate_fn)
    
    # Train
    train_loss = train(model, dataloader, criterion, optimizer, device)
    print(f'Training complete. Loss: {train_loss:.4f}')
    
    # Print a message indicating readiness for generating responses
    print("Model trained. Ready to generate responses.")
    print("You can now input prompts for response generation.")
    signal.signal(signal.SIGINT, signal.SIG_DFL)  # Restore default signal handler

# Register the signal handler
signal.signal(signal.SIGINT, handle_interrupt)

print("Input new training data. Press Ctrl+C to stop and train.")

# Interactive input loop
try:
    while True:
        src = input("Enter source sentence: ")
        trg = input("Enter target sentence: ")
        source_sentences.append(src)
        target_sentences.append(trg)
        
        # Update vocabularies
        source_vocab = build_vocab(source_sentences)
        target_vocab = build_vocab(target_sentences)
except KeyboardInterrupt:
    pass
cannon split text.split() because it is a list


# After training, generate responses
print("Ready for response generation.")
while True:
    prompt = input("Enter prompt: ")
    response = generate_response(prompt, model, source_vocab, target_vocab, device)
    print(f"Response: {response}")