Adityadn's picture
Update app.py
b21ab62 verified
raw
history blame
1.07 kB
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()