Update app.py

app.py

@@ -21,14 +21,18 @@ def main():
 import requests
 
 
+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()
 
-from transformers import pipeline
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
 
-pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
+# from transformers import pipeline
+
+# pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
 
 
 
@@ -39,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(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)}