import streamlit as st

def main():
    st.title("Amazon Title Suggestion")
    if "title" not in st.session_state:
        st.session_state.title = ""
    if "ner_dict" not in st.session_state:
        st.session_state.ner_dict = {}
    if "selected_keywords" not in st.session_state:
        st.session_state.selected_keywords = []
    if "submitted_title" not in st.session_state:
        st.session_state.submitted_title = False
    if "submitted_ner_keywords" not in st.session_state:
        st.session_state.submitted_ner_keywords = False
    
    if not st.session_state.submitted_title:
        submit_title()
    elif st.session_state.submitted_title and not st.session_state.submitted_ner_keywords:
        submit_ner_keywords()

import requests



# def query(payload):
#     response = requests.post(API_URL, headers=headers, json=payload)
#     return response.json()

from transformers import pipeline

pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")



def ner_title(title):
    # Define the roles and markers
    B_SYS, E_SYS = "<<SYS>>", "<</SYS>>"
    B_INST, E_INST = "[INST]", "[/INST]"
    B_in, E_in = "[Title]", "[/Title]"
    # Format your prompt template
    prompt = f"""{B_INST} {B_SYS} You are a helpful assistant that provides accurate and concise responses. {E_SYS}\nExtract named entities from the given product title. Provide the output in JSON format.\n{B_in} {title.strip()} {E_in}\n{E_INST}\n\n### NER Response:\n{{"{title.split()[0].lower()}"""
    # output = query({
    # "inputs": prompt,
    # })

    return eval(pipe(text)[0]["generated_text"].split("### NER Response:\n")[-1])
    #return(eval(output[0]['generated_text'].split("### NER Response:\n")[-1]))
# def ner_title(title):
#     word_list = title.split()
#     indexed_dict = {index: word for index, word in enumerate(word_list)}
#     return indexed_dict

def submit_title():
    title = st.text_input("Enter Product Title:")
    if st.button("Submit Title"):
        st.session_state.title = title
        ner = ner_title(title)
        st.session_state.submitted_title = True
        st.session_state.ner_dict = ner

def submit_ner_keywords():
    st.subheader("Product Features:")
    selected_features = []
    for key, value in st.session_state.ner_dict.items():
        if st.checkbox(f"{key}: {value}"):
            selected_features.append(value)
    
    st.subheader("Select Search Terms:")
    keyword_list = ['a','b','c','f','g',"Feature", "Price", "Quality", "Availability"]
    for keyword in keyword_list:
        st.checkbox(keyword, key=keyword)
    
    if st.button("Suggest Titles"):
        model2_keywords = [keyword for keyword in keyword_list if st.session_state[keyword]]
        st.session_state.selected_keywords = model2_keywords
        st.session_state.submitted_ner_keywords = True
        st.write("Selected Keywords for Model2:", model2_keywords)
        st.write("Selected features for Model2:", selected_features)
    
    if st.button("Reset"):
        st.session_state.title = ""
        st.session_state.submitted_title = False
        st.session_state.submitted_ner_keywords = False
        # Reset selected keywords
        for keyword in keyword_list:
            st.session_state[keyword] = False
        # Rerun the app
        st.experimental_rerun()

if __name__ == "__main__":
    main()