import torch import torch.nn as nn from torch.nn import functional as F import tiktoken import gradio as gr # Define the model architecture class GPTConfig: def __init__(self): self.block_size = 1024 self.vocab_size = 50304 self.n_layer = 12 self.n_head = 12 self.n_embd = 768 class CausalSelfAttention(nn.Module): def __init__(self, config): super().__init__() assert config.n_embd % config.n_head == 0 self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd) self.c_proj = nn.Linear(config.n_embd, config.n_embd) self.n_head = config.n_head self.n_embd = config.n_embd self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size)) def forward(self, x): B, T, C = x.size() q, k, v = self.c_attn(x).split(self.n_embd, dim=2) k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True) y = y.transpose(1, 2).contiguous().view(B, T, C) return self.c_proj(y) class MLP(nn.Module): def __init__(self, config): super().__init__() self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd) self.gelu = nn.GELU() self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd) def forward(self, x): return self.c_proj(self.gelu(self.c_fc(x))) class Block(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd) self.attn = CausalSelfAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd) self.mlp = MLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x class GPT(nn.Module): def __init__(self, config): super().__init__() self.config = config self.transformer = nn.ModuleDict(dict( wte = nn.Embedding(config.vocab_size, config.n_embd), wpe = nn.Embedding(config.block_size, config.n_embd), h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) self.transformer.wte.weight = self.lm_head.weight self.apply(self._init_weights) def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: torch.nn.init.zeros_(module.bias) elif isinstance(module, nn.Embedding): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, idx, targets=None): device = idx.device b, t = idx.size() assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0) tok_emb = self.transformer.wte(idx) pos_emb = self.transformer.wpe(pos) x = tok_emb + pos_emb for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) logits = self.lm_head(x) loss = None if targets is not None: loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1) return logits, loss # Load the model def load_model(model_path): config = GPTConfig() model = GPT(config) checkpoint = torch.load(model_path, map_location=torch.device('cpu')) print("Checkpoint keys:", checkpoint.keys()) # Debug print if 'model_state_dict' in checkpoint: model.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint) model.eval() return model # Load the model model = load_model('gpt_5000.pt') # Replace with the actual path to your .pt file enc = tiktoken.get_encoding('gpt2') # Improved text generation function def generate_text(prompt, max_length=100, temperature=0.7, top_k=50): input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0) generated = [] with torch.no_grad(): for _ in range(max_length): outputs, _ = model(input_ids) next_token_logits = outputs[:, -1, :] # Apply temperature next_token_logits = next_token_logits / temperature # Apply top-k filtering top_k_logits, top_k_indices = torch.topk(next_token_logits, top_k, dim=-1) next_token_probs = F.softmax(top_k_logits, dim=-1) # Sample from the filtered distribution next_token_index = torch.multinomial(next_token_probs, num_samples=1) next_token = top_k_indices.gather(-1, next_token_index) input_ids = torch.cat([input_ids, next_token], dim=-1) generated.append(next_token.item()) # Stop if we generate a newline, but only after generating at least 20 tokens if next_token.item() == enc.encode('\n')[0] and len(generated) > 20: break generated_text = enc.decode(generated) return prompt + generated_text # Gradio interface def gradio_generate(prompt, max_length, temperature, top_k): return generate_text(prompt, max_length, temperature, top_k) iface = gr.Interface( fn=gradio_generate, inputs=[ gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."), gr.Slider(minimum=20, maximum=500, value=100, step=1, label="Max Length"), gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature"), gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k") ], outputs=gr.Textbox(label="Generated Text"), title="GPT Text Generator", description="Enter a prompt and adjust parameters to generate text using a fine-tuned GPT model." ) # Launch the app iface.launch()