Create app.py

app.py

@@ -0,0 +1,72 @@
+import gradio as gr
+import torch
+from torch.nn import functional as F
+from gpt_class import GPTConfig, GPT
+import tiktoken
+
+# Setup device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model
+state_dict = torch.load('log/model_51999.pt', map_location=device)
+config = state_dict['config']
+model = GPT(config)
+model.load_state_dict(state_dict['model'])
+model.to(device)
+model.eval()
+
+# Set seed for reproducibility
+torch.manual_seed(42)
+torch.cuda.manual_seed_all(42)
+
+# Get tokenizer
+tokenizer = tiktoken.get_encoding("gpt2")
+
+def generate_text(example, num_return_sequences='4', max_length='64'):
+    num_return_sequences = int(num_return_sequences) if num_return_sequences.isdigit() else 4
+    max_length = int(max_length) if max_length.isdigit() else 64
+
+    model.eval()
+    tokens = tokenizer.encode(example)
+    tokens = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).repeat(num_return_sequences, 1)
+    tokens = tokens.to(device)
+    sample_rng = torch.Generator(device=device)
+
+    xgen = tokens
+    while xgen.size(1) < max_length:
+        with torch.no_grad():
+            with torch.autocast(device_type=device):
+                logits, _ = model(xgen)  # Assumes model returns logits and optional loss
+            logits = logits[:, -1, :]  # Get last token logits
+            probs = F.softmax(logits, dim=-1)
+            topk_probs, topk_indices = torch.topk(probs, 50, dim=-1)
+            ix = torch.multinomial(topk_probs, 1, generator=sample_rng)
+            xcol = torch.gather(topk_indices, -1, ix)
+            xgen = torch.cat((xgen, xcol), dim=1)
+
+    results = []
+    for i in range(num_return_sequences):
+        tokens = xgen[i, :max_length].tolist()
+        decoded = tokenizer.decode(tokens)
+        results.append(decoded)
+    return "\n\n".join(results)
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.components.Textbox(label="Prompt"),
+        gr.components.Textbox(label="Number of Sequences [1-4]"),
+        gr.components.Textbox(label="Maximum Length [32-128]")
+    ],
+    outputs=gr.components.Textbox(label="Generated Text"),
+    title="Text Generator",
+    description="Enter a prompt to generate text using a GPT model. Adjust the number of sequences and the maximum length as needed.",
+    examples=[
+        ["It is raining and my family", "2", "64"],
+        ["We entered into the forest and", "2", "64"],
+        ["I sat for doing my homework", "2", "64"]
+    ]
+)
+
+iface.launch(share=True)