token

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 from peft import PeftModel
 import spaces  # Ensure spaces is imported
@@ -8,14 +8,6 @@ import spaces  # Ensure spaces is imported
 # Define the base model ID
 base_model_id = "meta-llama/Llama-2-13b-hf"
 
-# Create a BitsAndBytesConfig object with the corrected settings
-quantization_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_use_double_quant=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
-
 # Ensure you have the Hugging Face token set as an environment variable
 huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 if not huggingface_token:
@@ -24,14 +16,13 @@ if not huggingface_token:
 # Define cache directory
 cache_dir = "./cache"
 
-# Load the base model with the updated quantization configuration
+# Load the base model without quantization configuration
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_id,  
-    quantization_config=quantization_config,
     trust_remote_code=True,
     token=huggingface_token,  # Use the token parameter
     cache_dir=cache_dir  # Specify cache directory
-)
+).to("cuda")  # Move model to CUDA
 
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
@@ -43,7 +34,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 # Load the fine-tuned model
-ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir)
+ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir).to("cuda")  # Move model to CUDA
 
 def formatting_func(job_description):
     text = f"### The job description: {job_description}\n ### The skills: "