import gradio as gr

#gr.Interface.load("models/nitinbhayana/Llama-2-7b-chat-hf-review-phrases-sentiments-v2").launch()


from transformers import pipeline

#pipeline = pipeline("text-generation", model="nitinbhayana/Llama-2-7b-chat-hf-review-phrases-sentiments-v2")


pipeline = pipeline("text-generation", model="nitinbhayana/TinyLlama-1.1B-Chat-v1.0-title-suggestion-v1.0")

def predict(title):
    
    prompt=f"""<s>[INST] <<SYS>>
You are a helpful, respectful, and honest assistant for product feature mapping.
<</SYS>>

[INST]
For a given e-commerce product title.
[TITLE]
Pantene Hair Science Silky Smoth Shampo with Pro-Vitamins & Vitamin E, shampoo for woman & men, shampoo for frizzy and dry hairs, 650ml
[/TITLE]
Create a new, easy-to-read, error-free title.
[/INST] 

Suggested Title:
"""
    predictions = pipeline(prompt)
    return (predictions)

gr.Interface(
    predict,
    inputs='text',
    outputs='text',
    title="Title Suggestion",
).launch()