Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 # Replace with your model name
 #MODEL_NAME = "ssirikon/Gemma7b-bnb-Unsloth"
 #MODEL_NAME = "unsloth/gemma-7b-bnb-4bit"
-MODEL_NAME = "unsloth/mistral-7b-bnb-4bit"
+MODEL_NAME = "Lohith9459/gemma7b"
 
 # Load the model and tokenizer
 model = AutoModelForCausalLM.from_pretrained(
@@ -17,33 +17,40 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# **Change 1: Set `llm_int8_skip_modules` to avoid deep copy**
-#model.quantization_config.llm_int8_skip_modules = ['lm_head'] 
-
-# Create a pipeline for text generation
-generator = pipeline(
-    task="summarization",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=50,  # Adjust as needed
-    do_sample=True,
-    top_k=10,
-    num_return_sequences=1,
-    eos_token_id=tokenizer.eos_token_id,
-)
-
-def generate_text(email):
-    result = generator("Generate a subject line for the following email.\n"+email)
-    return result[0]["generated_text"]
-
 
-# Create a Gradio interface
+def generate_subject(email_body):
+  instruction = "Generate a subject line for the following email."
+  formatted_text = f"""Below is an instruction that describes a task. \
+    Write a response that appropriately completes the request.
+
+    ### Instruction:
+    {instruction}
+
+    ### Input:
+    {email_body}
+
+    ### Response:
+    """
+  inputs = tokenizer([formatted_text], return_tensors="pt").to("cuda")
+  text_streamer = TextStreamer(tokenizer)
+  generated_ids = model.generate(**inputs, streamer=text_streamer, max_new_tokens=512)
+  generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+  
+  def extract_subject(text):
+    start_tag = "### Response:"
+    start_idx = text.find(start_tag)
+    if start_idx == -1:
+        return None
+    subject = text[start_idx + len(start_tag):].strip()
+    return subject
+  
+  return extract_subject(generated_text)
+
+# Create the Gradio interface
 demo = gr.Interface(
-    fn=generate_text,
-    inputs=gr.Textbox(lines=5, label="Enter your Email here:"),
-    outputs=gr.Textbox(label="Generated Subject"),
-    title="Email Subject Generation demo",
-    description="Enter an email and let the model generate the subject for you!",
+    fn=generate_subject,
+    inputs=gr.Textbox(lines=20, label="Email Body"),
+    outputs=gr.Textbox(label="Generated Subject")
 )
 
-demo.launch(debug=True)
+demo.launch()