File size: 2,021 Bytes
dbb4d61 59fc6f3 dbb4d61 95deaf9 dbb4d61 42884c7 dbb4d61 42884c7 dbb4d61 |
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 |
import gradio as gr
import torch
import tiktoken
import numpy as np
from model import GPT, GPTConfig # Changed from train to model
def load_quantized_model():
model = GPT(GPTConfig())
quantized_dict = torch.load("gpt_model_quantized.pt")
# Dequantize model
state_dict = {}
for key, value in quantized_dict.items():
if isinstance(value, dict):
state_dict[key] = torch.tensor(
value['data'].astype(np.float32) * value['scale']
)
else:
state_dict[key] = value
model.load_state_dict(state_dict)
model.eval()
return model
def generate_text(input_text):
try:
# Set device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Load model
model = load_quantized_model()
model = model.to(device)
# Tokenize input
tokenizer = tiktoken.get_encoding('gpt2')
input_tokens = torch.tensor([tokenizer.encode(input_text)]).to(device)
# Generate
with torch.no_grad():
output_tokens = model.generate(input_tokens, max_new_tokens=10)[0].tolist()
# Decode and return
generated_text = tokenizer.decode(output_tokens)
return generated_text
except Exception as e:
return f"Error generating text: {e}"
# Create Gradio interface
iface = gr.Interface(
fn=generate_text,
inputs=gr.Textbox(lines=5, label="Input Text"),
outputs=gr.Textbox(lines=10, label="Generated Text"),
title="Text Generator",
description="Enter some text and the model will generate a Shakespeare-style continuation.",
examples=[
["To be, or not to be,"],
["All the world's a stage, and all the men"],
["But soft, what light through yonder"],
["Friends, Romans, countrymen,"],
["Now is the winter of our discontent"]
]
)
# Launch the interface
iface.launch() |