File size: 3,410 Bytes
26a4eca
458a9b1
26a4eca
 
 
 
 
 
 
 
 
 
 
 
458a9b1
26a4eca
 
458a9b1
26a4eca
458a9b1
3025fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
26a4eca
3025fbf
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
458a9b1
26a4eca
 
 
 
458a9b1
26a4eca
458a9b1
 
26a4eca
 
 
 
 
 
458a9b1
26a4eca
c02042b
26a4eca
 
458a9b1
26a4eca
 
 
 
 
 
458a9b1
26a4eca
 
 
 
c02042b
 
 
 
458a9b1
 
26a4eca
458a9b1
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
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()