Update app.py

app.py

@@ -18,10 +18,37 @@ def main():
     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):
-    word_list = title.split()
-    indexed_dict = {index: word for index, word in enumerate(word_list)}
-    return indexed_dict
+    # 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:")