Update app.py

app.py

@@ -1,51 +1,55 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from unsloth import FastLanguageModel
+from transformers import TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Replace with your model name
-#MODEL_NAME = "ssirikon/Gemma7b-bnb-Unsloth"
+MODEL_NAME = "ssirikon/Gemma7b-bnb-Unsloth"
 #MODEL_NAME = "unsloth/gemma-7b-bnb-4bit"
-MODEL_NAME = "google-t5/t5-small"
-#MODEL_NAME = "unsloth/mistral-7b-bnb-4bit"
+#MODEL_NAME = "Lohith9459/gemma7b"
 
 # Load the model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, 
-    device_map="auto", 
-    torch_dtype=torch.float16,
-    load_in_4bit=True,  # Load the model in 4-bit precision
-    # Removed the unsupported argument
-)
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+max_seq_length = 512
+dtype = torch.bfloat16
+load_in_4bit = True
 
-# **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="text-generation",
-    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"]
+#model = FastLanguageModel.from_pretrained(MODEL_NAME, max_seq_length=max_seq_length, dtype=dtype, load_in_4bit=load_in_4bit)
+#tokenizer = model.tokenizer
 
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# 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()