Spaces:
Runtime error
Runtime error
import gradio as gr | |
from transformers import DistilGPT2LMHeadModel, DistilGPT2Tokenizer | |
# Load pre-trained DistilGPT-2 model and tokenizer | |
tokenizer = DistilGPT2Tokenizer.from_pretrained("distilgpt2") | |
model = DistilGPT2LMHeadModel.from_pretrained("distilgpt2") | |
# Function to generate keywords from input text | |
def generate_keywords(prompt): | |
# Encode the input text | |
inputs = tokenizer.encode(prompt, return_tensors="pt") | |
# Generate output | |
outputs = model.generate(inputs, max_length=50, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95) | |
# Decode and return the generated text (keywords) | |
keywords = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return keywords.strip() | |
# Gradio interface | |
iface = gr.Interface(fn=generate_keywords, | |
inputs=gr.Textbox(label="Enter Ad Prompt", placeholder="E.g., Generate ad keywords for wireless headphones"), | |
outputs=gr.Textbox(label="Generated Keywords"), | |
live=True) | |
# Launch the Gradio interface | |
iface.launch() | |