Update app.py

app.py

@@ -19,20 +19,20 @@ def main():
         submit_ner_keywords()
 
 import requests
-import os
+# import os
 
-API_URL = "https://api-inference.huggingface.co/models/shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0"
-TOKEN = os.environ.get("HUGGING_FACE_TOKEN")
-headers = {"Authorization": f"Bearer {TOKEN}"}
+# API_URL = "https://api-inference.huggingface.co/models/shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0"
+# TOKEN = os.environ.get("HUGGING_FACE_TOKEN")
+# headers = {"Authorization": f"Bearer {TOKEN}"}
 
 
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+# def query(payload):
+#     response = requests.post(API_URL, headers=headers, json=payload)
+#     return response.json()
 
-# from transformers import pipeline
+from transformers import pipeline
 
-# pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
+pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
 
 
 
@@ -43,12 +43,12 @@ def ner_title(title):
     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,
-    })
+    # 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]))
+    return eval(pipe(prompt)[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)}