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

1import torch
2import numpy as np
3
4# Configuration (MUST match training configuration)
5CONFIG = {
6    "FILE_PATH": 'dataset.txt',
7    "SEQ_LENGTH": 32,
8    "EMBEDDING_DIM": 64,
9    "HIDDEN_DIM": 64,
10    "NUM_LAYERS": 1,
11    "DROPOUT": 0.2,
12    "MODEL_SAVE_PATH": "char_lm_advanced.pth",
13    "TEMPERATURE": 0.7,
14    "TOP_K": 5,
15    "TOP_P": 0.95
16}
17
18# Load vocabulary
19with open(CONFIG["FILE_PATH"], 'r', encoding='utf-8') as f:
20    text = f.read()
21
22chars = sorted(list(set(text)))
23char_to_idx = {ch: i for i, ch in enumerate(chars)}
24idx_to_char = {i: ch for i, ch in enumerate(chars)}
25vocab_size = len(chars)
26
27# Model definition (must match training architecture)
28class CharLM(torch.nn.Module):
29    def __init__(self):
30        super(CharLM, self).__init__()
31        self.embedding = torch.nn.Embedding(vocab_size, CONFIG["EMBEDDING_DIM"])
32        self.lstm = torch.nn.LSTM(
33            CONFIG["EMBEDDING_DIM"], 
34            CONFIG["HIDDEN_DIM"],
35            num_layers=CONFIG["NUM_LAYERS"],
36            dropout=CONFIG["DROPOUT"] if CONFIG["NUM_LAYERS"] > 1 else 0,
37            batch_first=True
38        )
39        self.dropout = torch.nn.Dropout(CONFIG["DROPOUT"])
40        self.fc = torch.nn.Linear(CONFIG["HIDDEN_DIM"], vocab_size)
41        
42    def forward(self, x, hidden=None):
43        x = self.embedding(x)
44        out, hidden = self.lstm(x, hidden)
45        out = self.dropout(out)
46        out = self.fc(out)
47        return out, hidden
48
49# Load trained model
50model = CharLM()
51model.load_state_dict(torch.load(CONFIG["MODEL_SAVE_PATH"]))
52model.eval()
53
54def generate_text(model, start_str, length=200, temperature=CONFIG["TEMPERATURE"], 
55                 top_k=CONFIG["TOP_K"], top_p=CONFIG["TOP_P"]):
56    """
57    Generate text with temperature scaling, top-k, and nucleus (top-p) sampling
58    """
59    model.eval()
60    chars = list(start_str)
61    input_seq = torch.tensor([char_to_idx[ch] for ch in chars]).unsqueeze(0)
62    hidden = None
63    
64    with torch.no_grad():
65        for _ in range(length):
66            outputs, hidden = model(input_seq, hidden)
67            logits = outputs[0, -1] / temperature
68            
69            # Apply top-k filtering
70            if top_k > 0:
71                top_vals, top_idx = torch.topk(logits, top_k)
72                logits[logits < top_vals[-1]] = -float('Inf')
73            
74            # Apply nucleus (top-p) filtering
75            if top_p > 0:
76                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
77                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
78                sorted_indices_to_remove = cumulative_probs > top_p
79                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
80                sorted_indices_to_remove[..., 0] = 0
81                indices_to_remove = sorted_indices[sorted_indices_to_remove]
82                logits[indices_to_remove] = -float('Inf')
83
84            probs = torch.softmax(logits, dim=-1)
85            next_char = torch.multinomial(probs, num_samples=1).item()
86            chars.append(idx_to_char[next_char])
87            input_seq = torch.tensor([[next_char]])
88    
89    return ''.join(chars)
90
91# Interactive loop
92while True:
93    try:
94        print("\n" + "="*50)
95        prompt = input("Enter your starting text (or 'exit' to quit):\n> ")
96        
97        if prompt.lower() == 'exit':
98            print("Goodbye!")
99            break
100            
101        # Filter invalid characters
102        valid_prompt = [c for c in prompt if c in char_to_idx]
103        if not valid_prompt:
104            print("Please use characters from the training data.")
105            continue
106            
107        # Get generation parameters
108        length = int(input("Output length (50-500 recommended): ")) or 200
109        temp = float(input(f"Temperature [{CONFIG['TEMPERATURE']}]: ") or CONFIG["TEMPERATURE"])
110        top_k = int(input(f"Top-K [{CONFIG['TOP_K']}]: ") or CONFIG["TOP_K"])
111        top_p = float(input(f"Top-P [{CONFIG['TOP_P']}]: ") or CONFIG["TOP_P"])
112        
113        # Generate and display
114        print("\nGenerating...")
115        generated = generate_text(
116            model,
117            ''.join(valid_prompt),
118            length=length,
119            temperature=temp,
120            top_k=top_k,
121            top_p=top_p
122        )
123        print("\nGenerated Text:")
124        print(generated)
125        print("="*50)
126        
127    except ValueError:
128        print("Invalid input! Please enter valid numbers for parameters.")
129    except KeyboardInterrupt:
130        print("\nExiting...")
131        break