Update app.py

app.py

@@ -1,20 +1,19 @@
 import pandas as pd
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 import os
-
+bnb_config = BitsAndBytesConfig(load_in_4bit=True)
 model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
 hf_token = os.environ.get("HF_TOKEN")
 
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
-
-# model = AutoModelForCausalLM.from_pretrained(
-#     model_name,
-#     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-#     device_map="auto"
-# )
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto",
+    use_auth_token=hf_token
+)
 
 def generate_prompt(original, translation):
     return f"### Task: Machine Translation Quality Estimation\n\nSource: {original}\nTranslation: {translation}\n\nScore (0-1):"