Update app.py

app.py

@@ -1,35 +1,27 @@
 import gradio as gr
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load pre-trained model and tokenizer
-model_name = "gpt2"  # You can use other models like gpt-2-large or gpt-3 for better performance
-model = GPT2LMHeadModel.from_pretrained(model_name)
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+# Load the pre-trained LLaMA model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("facebook/llama-7b")
+model = AutoModelForCausalLM.from_pretrained("facebook/llama-7b")
 
-# Function to generate keywords based on a prompt
-def generate_keywords(prompt):
-    # Encode input prompt with a more direct instruction for only keywords
-    prompt_with_instruction = prompt + " Only provide a list of keywords, no additional text."
-    inputs = tokenizer.encode(prompt_with_instruction, return_tensors="pt")
+# Function to generate keywords from input text
+def generate_keywords(text):
+    # Encode the input text
+    inputs = tokenizer.encode(text, return_tensors="pt")
 
-    # Generate output from model
+    # Generate the output from the model
     outputs = model.generate(inputs, max_length=50, num_return_sequences=1, no_repeat_ngram_size=2, top_k=50, top_p=0.95)
 
-    # Decode generated tokens
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Clean up the text to remove unnecessary parts
-    # Remove anything after 'Only provide a list of keywords'
-    clean_text = generated_text.split("Only provide a list of keywords")[0].strip()
-    
-    # Return the keywords only
-    return clean_text
+    # Decode and return the generated 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"),
+                     inputs=gr.Textbox(label="Enter Prompt", placeholder="E.g., Generate ad keywords for wireless headphones"), 
                      outputs=gr.Textbox(label="Generated Keywords"),
                      live=True)
 
-# Launch interface
+# Launch the Gradio interface
 iface.launch()