token

app.py

@@ -2,7 +2,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
 from peft import PeftModel
-import spaces
+import os
 
 # Define the base model ID
 base_model_id = "meta-llama/Llama-2-13b-hf"
@@ -15,17 +15,21 @@ quantization_config = BitsAndBytesConfig(
     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:
+    raise Exception("Hugging Face token not found. Please set it as an environment variable 'HUGGINGFACE_TOKEN'.")
+
 # Load the base model with the updated quantization configuration
-# Adjust 'device_map' based on your system's GPU configuration
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_id,  
     quantization_config=quantization_config,
     trust_remote_code=True,
-    token=True  # Update this to use the token parameter
+    token=huggingface_token  # Use the token parameter
 )
 
 # Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(base_model_id, add_bos_token=True, trust_remote_code=True, token=huggingface_token)
 
 # Load the fine-tuned model
 ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800")
@@ -34,7 +38,6 @@ def formatting_func(job_description):
     text = f"### The job description: {job_description}\n ### The skills: "
     return text
 
-@spaces.GPU
 def generate_skills(job_description):
     formatted_text = formatting_func(job_description)
     model_input = tokenizer(formatted_text, return_tensors="pt").to("cuda")