Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConfig
+from unsloth import FastLanguageModel
+from transformers import TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Replace with your model name
 #MODEL_NAME = "ssirikon/Gemma7b-bnb-Unsloth"
@@ -8,35 +10,28 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, AutoConf
 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)
-config = AutoConfig.from_pretrained(MODEL_NAME)
+max_seq_length = 512
+dtype = torch.bfloat16
+load_in_4bit = True
 
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 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)
@@ -44,7 +39,7 @@ def generate_subject(email_body):
         return None
     subject = text[start_idx + len(start_tag):].strip()
     return subject
-  
+
   return extract_subject(generated_text)
 
 # Create the Gradio interface