Update app.py

app.py

@@ -1,25 +1,25 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import DistilGPT2LMHeadModel, DistilGPT2Tokenizer
 
-# Load the pre-trained LLaMA model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("facebook/llama-7b")
-model = AutoModelForCausalLM.from_pretrained("facebook/llama-7b")
+# 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(text):
+def generate_keywords(prompt):
     # Encode the input text
-    inputs = tokenizer.encode(text, return_tensors="pt")
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
 
-    # Generate the output from the model
+    # 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 keywords
+    # 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 Prompt", placeholder="E.g., Generate ad keywords for wireless headphones"), 
+                     inputs=gr.Textbox(label="Enter Ad Prompt", placeholder="E.g., Generate ad keywords for wireless headphones"), 
                      outputs=gr.Textbox(label="Generated Keywords"),
                      live=True)