File size: 1,329 Bytes
aa60148 |
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 |
import gradio as gr
import torch
from bigram_model import BigramLanguageModel, encode, decode
# Assuming 'BigramLanguageModel' and 'decode' are defined as in your code
class GradioInterface:
def __init__(self, model_path=None):
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.model = self.load_model(model_path)
self.model.eval()
def load_model(self, model_path):
model = BigramLanguageModel().to(self.device)
if model_path:
model.load_state_dict(torch.load(model_path, map_location=self.device))
return model
def generate_text(self, input_text, max_tokens=100):
context = torch.tensor([encode(input_text)], dtype=torch.long, device=self.device)
output = self.model.generate(context, max_new_tokens=max_tokens)
return decode(output[0].tolist())
# Load the model
model_path = "models/lafontaine_gpt_v8_241011_1307.pth"
model_interface = GradioInterface(model_path)
# Define Gradio interface
gr_interface = gr.Interface(
fn=model_interface.generate_text,
inputs=["text", gr.Slider(50, 500)],
outputs="text",
description="Bigram Language Model text generation. Enter some text, and the model will continue it.",
examples=[["Once upon a time"]]
)
# Launch the interface
gr_interface.launch()
|