ysn-rfd/text-dataset-tiny-code-script-py-format

1import torch
2import torch.nn as nn
3
4# Configuration
5MODEL_SAVE_PATH = "char_lm_model.pth"
6SEQ_LENGTH = 32
7EMBEDDING_DIM = 64
8HIDDEN_DIM = 64
9
10# Load vocabulary
11with open('dataset.txt', 'r', encoding='utf-8') as f:
12    text = f.read()
13
14chars = sorted(list(set(text)))
15vocab_size = len(chars)
16char_to_idx = {ch: i for i, ch in enumerate(chars)}
17idx_to_char = {i: ch for i, ch in enumerate(chars)}
18
19# Model architecture
20class CharLM(nn.Module):
21    def __init__(self):
22        super(CharLM, self).__init__()
23        self.embedding = nn.Embedding(vocab_size, EMBEDDING_DIM)
24        self.rnn = nn.GRU(EMBEDDING_DIM, HIDDEN_DIM, batch_first=True)
25        self.fc = nn.Linear(HIDDEN_DIM, vocab_size)
26        
27    def forward(self, x, hidden=None):
28        x = self.embedding(x)
29        out, hidden = self.rnn(x, hidden)
30        out = self.fc(out)
31        return out, hidden
32
33# Load the trained model
34model = CharLM()
35model.load_state_dict(torch.load(MODEL_SAVE_PATH))
36model.eval()
37
38def generate_text(model, start_str, length=100, temperature=0.7, top_k=0):
39    """
40    Generate text with temperature scaling and top-k sampling
41    """
42    model.eval()
43    chars = [ch for ch in start_str]
44    input_seq = torch.tensor([char_to_idx[ch] for ch in chars]).unsqueeze(0)
45    hidden = None
46    
47    with torch.no_grad():
48        for _ in range(length):
49            outputs, hidden = model(input_seq, hidden)
50            logits = outputs[0, -1] / temperature
51            
52            if top_k > 0:
53                top_vals, top_idx = torch.topk(logits, top_k)
54                logits[logits < top_vals[-1]] = -float('Inf')
55            
56            probs = torch.softmax(logits, dim=-1)
57            next_char = torch.multinomial(probs, num_samples=1).item()
58            chars.append(idx_to_char[next_char])
59            input_seq = torch.tensor([[next_char]])
60    
61    return ''.join(chars)
62
63# Chat loop
64def chat():
65    print("Chat with the model! Type 'exit' to stop.")
66    while True:
67        user_input = input("You: ")
68        if user_input.lower() == 'exit':
69            break
70        response = generate_text(model, user_input, length=100, temperature=0.7, top_k=5)
71        print("Bot:", response)
72
73chat()
74