Spaces:
Build error
Build error
File size: 2,495 Bytes
fa9c2df 51a0934 fa9c2df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 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('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) |