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 assistant that provides accurate and concise responses. <</SYS>>
Create a new, easy to read, and error free title for a given Ecommerce product title.
[Title] {title} [/Title]
[/INST]
### Suggested 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]
# {title}
# [/TITLE]
# Create a new, easy-to-read, error-free title.
# [/INST] 

# Suggested Title:
# """
    predictions = pipeline(prompt)
    output=((predictions[0]['generated_text']).split('Suggested Title:')[-1]).strip()
    return (output)

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