Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 # Load the Hugging Face API token from environment variable
 token = os.getenv("HUGGINGFACE_API_TOKEN")
@@ -12,19 +12,26 @@ if not token:
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the tokenizer and model using the token
-model_repo = "unsloth/Llama-3.2-3B-Instruct-bnb-4bit"
-tokenizer = AutoTokenizer.from_pretrained(model_repo, use_auth_token=token)
+model_repo = "unsloth/llama-3.2-3b-instruct-bnb-4bit"
+tokenizer = AutoTokenizer.from_pretrained(model_repo, token=token)
+
+# Configure 4-bit quantization
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16
+)
+
+# Load the model with quantization configuration
 model = AutoModelForCausalLM.from_pretrained(
     model_repo,
-    use_auth_token=token,
+    token=token,
     device_map="auto",
-    torch_dtype=torch.float16,
-    load_in_4bit=True,
-    quantization_config={"bnb_4bit_use_double_quant": True, "bnb_4bit_quant_type": "nf4"}
+    quantization_config=quantization_config
 )
 
-# Move the model to the device
-model.to(device)
+# Ensure the model is in evaluation mode
 model.eval()
 
 # Define the inference function
@@ -37,7 +44,7 @@ def infer(prompt):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## Llama 3.2 3B Instruct Model Inference")
+    gr.Markdown("## LLaMA 3.2 3B Instruct Model Inference")
 
     with gr.Row():
         prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here...")