Update app.py

app.py

@@ -2,10 +2,12 @@ import gradio as gr
 
 # gr.load("models/shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0").launch()
 # Load model directly
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-tokenizer = AutoTokenizer.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0")
-model = AutoModelForCausalLM.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0")
+pipe = pipeline("text-generation", model="shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0")
+
+# tokenizer = AutoTokenizer.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0")
+# model = AutoModelForCausalLM.from_pretrained("shivanikerai/TinyLlama-1.1B-Chat-v1.0-seo-optimised-title-suggestion-v1.0")
 def generate_title_suggestions(keywords, product_info):
     # Define the roles and markers
     B_SYS, E_SYS = "<<SYS>>", "<</SYS>>"
@@ -17,26 +19,28 @@ def generate_title_suggestions(keywords, product_info):
     prompt = f"""{B_INST} {B_SYS} You are a helpful, respectful and honest assistant for ecommerce product title creation. {E_SYS}\nCreate a SEO optimized e-commerce product title for the keywords:{keywords.strip()}\n{B_in}{product_info}{E_in}\n{E_INST}\n\n{B_out}"""
 
 
-    print("Prompt:")
-    print(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=1024,
-                            do_sample=True,
-                            temperature=0.01,
-                            eos_token_id=tokenizer.eos_token_id,
-                            top_k=0)
+    # print("Prompt:")
+    # print(prompt)
+    predictions = pipeline(prompt)
+    output=((predictions[0]['generated_text']).split(B_out)[-1]).strip()
+    return (output)
+    # 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=1024,
+    #                         do_sample=True,
+    #                         temperature=0.01,
+    #                         eos_token_id=tokenizer.eos_token_id,
+    #                         top_k=0)
 
-    print()
+    # print()
 
     # 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_text = re.sub('\n+', '\n', output_text)  # remove excessive newline characters
+    # output_text = tokenizer.decode(output[0, len(encoding.input_ids[0]):], skip_special_tokens=False)
+    # output_text = re.sub('\n+', '\n', output_text)  # remove excessive newline characters
 
-    print("Generated Assistant Response:")
-    print(output_text)
+    # print("Generated Assistant Response:")
+    # print(output_text)
     gr.Interface(
     generate_title_suggestions,
     inputs='text',