Update app.py

app.py

@@ -1,16 +1,23 @@
 import streamlit as st
+import re
 
-import requests
-import os
-token=os.environ.get("HUGGING_FACE_TOKEN")
+# Use a pipeline as a high-level helper
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-API_URL = "https://api-inference.huggingface.co/models/shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0"
-headers = {"Authorization": f"Bearer {token}"}
+tokenizer = AutoTokenizer.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
+model = AutoModelForCausalLM.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0")
+# import requests
+# import os
+# token=os.environ.get("HUGGING_FACE_TOKEN")
 
+# API_URL = "https://api-inference.huggingface.co/models/shivanikerai/TinyLlama-1.1B-Chat-v1.0-sku-title-ner-generation-reversed-v1.0"
+# 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()
 
 def convert_to_dictionary(input_string):
     try:
@@ -117,9 +124,26 @@ 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,
+    # })
+
+    encoding = tokenizer(prompt, return_tensors="pt").to("cuda:0")
+    output = model.generate(input_ids=encoding.input_ids,
+                            attention_mask=encoding.attention_mask,
+                            max_new_tokens=512,
+                            do_sample=True,
+                            temperature=0.01,
+                            eos_token_id=tokenizer.eos_token_id,
+                            top_k=0)
+
+    
+
+    # Subtract the length of input_ids from output to get only the model's response
+    output_text = tokenizer.decode(output[0, len(encoding.input_ids[0]):], skip_special_tokens=False)
+    output = re.sub('\n+', '\n', output_text)  # remove excessive newline characters
+    #output = f"""{{\"{title.split()[0].lower()} {output_text}"""
+    #output = re.sub(' ": "', '": "', output)
 
     #return eval(pipe(text)[0]["generated_text"].split("### NER Response:\n")[-1])
     return convert_to_dictionary((output[0]['generated_text'].split("### NER Response:\n")[-1]))