import torch import torch.nn as nn from torch.nn import functional as F import tiktoken import gradio as gr try: import spaces use_spaces_gpu = True except ImportError: use_spaces_gpu = False def dummy_gpu_decorator(func): return func spaces = type('', (), {'GPU': dummy_gpu_decorator})() # Define the GPTConfig class 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 # Define other necessary classes 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 # Update the load_model function @spaces.GPU def load_model(model_path): config = GPTConfig() model = GPT(config) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") checkpoint = torch.load(model_path, map_location=device) if 'model_state_dict' in checkpoint: model.load_state_dict(checkpoint['model_state_dict']) else: model.load_state_dict(checkpoint) model.eval() model.to(device) return model enc = tiktoken.get_encoding('gpt2') # Update the generate_text function @spaces.GPU def generate_text(prompt, max_length=432, temperature=0.8, top_k=40): model = load_model('gpt_model.pth') device = next(model.parameters()).device input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0).to(device) generated = [] with torch.no_grad(): for _ in range(max_length): outputs, _ = model(input_ids) next_token_logits = outputs[:, -1, :] next_token_logits = next_token_logits / temperature 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) 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()) if next_token.item() == enc.encode('\n')[0] and len(generated) > 100: break return enc.decode(generated) # Add the gradio_generate function @spaces.GPU def gradio_generate(prompt, max_length, temperature, top_k): return generate_text(prompt, max_length, temperature, top_k) # # Your existing imports and model code here... css = """ """ with gr.Blocks(css=css) as demo: gr.HTML("

🌟 Enchanted Tales Generator 🌟

") with gr.Row(): with gr.Column(scale=3): prompt = gr.Textbox( placeholder="Begin your magical journey here (e.g., 'In a realm beyond the mists of time...')", label="Story Incantation", elem_classes="user-input" ) with gr.Column(scale=1): generate_btn = gr.Button("Weave the Tale", elem_classes="generate-btn") with gr.Row(): max_length = gr.Slider(minimum=50, maximum=500, value=432, step=1, label="Scroll Length") temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.8, step=0.1, label="Magical Intensity") top_k = gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Arcane Diversity") output = gr.Markdown(elem_classes="output-box") generate_btn.click( gradio_generate, inputs=[prompt, max_length, temperature, top_k], outputs=output ) gr.HTML("""
"In the realm of imagination, every word is a spell, every sentence a charm."
""") if __name__ == "__main__": demo.launch()