shahiil's picture
Create app.py
cf6f59d verified
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
import gradio as gr
# Load pre-trained Hugging Face model
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0" # Replace with your model
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=True)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=True)
# Initialize text generation pipeline
text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
def laptop_recommendation(user_input, task):
"""
Handles laptop recommendation tasks based on user preferences.
"""
if not user_input.strip():
return "Please provide some input."
if task == "Recommendation":
prompt = f"Recommend a laptop based on the following preferences:\n{user_input}\nRecommended Laptop:"
elif task == "Compare":
prompt = f"Compare two laptops based on the following specifications:\n{user_input}\nComparison:"
elif task == "Budget Recommendation":
prompt = f"Recommend the best laptop for the following budget:\n{user_input}\nRecommended Laptop for Budget:"
else:
return "Invalid task selected."
try:
response = text_generator(
prompt,
max_length=96,
num_return_sequences=1,
pad_token_id=tokenizer.eos_token_id,
temperature=0.7,
top_p=0.9
)[0]["generated_text"]
return response[len(prompt):].strip()
except Exception as e:
return f"An error occurred during text generation: {str(e)}"
def gradio_interface(user_input, task):
"""Gradio interface function."""
return laptop_recommendation(user_input, task)
with gr.Blocks() as laptop_recommendation_ui:
gr.Markdown("# Laptop Recommendation Chatbot")
gr.Markdown(
"This chatbot helps with recommending laptops based on preferences, comparing laptops, and suggesting options based on budget."
)
user_input = gr.Textbox(lines=5, placeholder="Enter your laptop preferences here...", label="Your Input")
task = gr.Radio(["Recommendation", "Compare", "Budget Recommendation"], label="Select Task")
output = gr.Textbox(lines=10, label="Chatbot Response")
submit_button = gr.Button("Submit")
clear_button = gr.Button("Clear")
submit_button.click(gradio_interface, inputs=[user_input, task], outputs=output)
clear_button.click(lambda: ("", ""), None, [user_input, output])
laptop_recommendation_ui.launch()