Spaces:
Sleeping
Sleeping
File size: 1,342 Bytes
6f14ee7 1390712 6f14ee7 1390712 |
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 |
import gradio as gr
import torch
from transformers import GPT2Tokenizer
# Assuming 'GPTLanguageModel' is already defined
class GPTLanguageModel(torch.nn.Module):
def forward(self, input_ids):
pass
def generate(self, input_ids, max_length=100):
return torch.tensor([[input_ids]]) # This is a placeholder for generation
class CustomTextGenerationPipeline:
def __init__(self, model, tokenizer):
self.model = model
self.tokenizer = tokenizer
def __call__(self, prompt, max_length=100):
input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
generated_ids = self.model.generate(input_ids, max_length=max_length)
return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
# Load model and tokenizer
model = GPTLanguageModel()
model.load_state_dict(torch.load("model.pth")) # Load the weights
model.eval()
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
# Create the pipeline
pipeline = CustomTextGenerationPipeline(model, tokenizer)
# Define the Gradio response function
def respond(message):
return pipeline(message, max_length=100)
# Create the Gradio interface
demo = gr.Interface(
fn=respond,
inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."),
outputs="text",
)
if __name__ == "__main__":
demo.launch()
|