Update app.py

app.py

@@ -1,24 +1,19 @@
-import gradio as gr
-import torch
+mport torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+import gradio as gr
 
-# Load base model without quantization for CPU compatibility
 base_model = AutoModelForCausalLM.from_pretrained(
-    "meta-llama/Llama-3.1-8B-Instruct",  # Use the original non-quantized model
+    "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
     torch_dtype=torch.float16,
-    device_map="cpu",  # Force CPU usage
-    low_cpu_mem_usage=True  # Optimize for CPU memory usage
+    device_map="auto",
+    load_in_4bit=True
 )
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
-
-# Add padding token if it doesn't exist
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+# tokenizer
+tokenizer = AutoTokenizer.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
 
-# Load LoRA adapters
+# LoRA adaptors
 model = PeftModel.from_pretrained(base_model, "rezaenayati/RezAi-Model")
 
 def chat_with_rezAi(messages, history):
@@ -30,19 +25,18 @@ def chat_with_rezAi(messages, history):
     
     conversation += f"<|start_header_id|>user<|end_header_id|>\n{messages}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n"
     
-    inputs = tokenizer([conversation], return_tensors="pt", padding=True, truncation=True, max_length=2048)
+    inputs = tokenizer([conversation], return_tensors="pt")
     
     with torch.no_grad():
         outputs = model.generate(
-            **inputs,
+            inputs,
             max_new_tokens=128,
             temperature=0.5,
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            attention_mask=inputs['attention_mask']
+            pad_token_id=tokenizer.eos_token_id
         )
     
-    # Get response
+    # get response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     new_response = response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()