Spaces:
Sleeping
Sleeping
| 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() | |