token

app.py

@@ -1,4 +1,4 @@
-import os  # Ensure os is imported
+import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
@@ -31,7 +31,7 @@ base_model = AutoModelForCausalLM.from_pretrained(
     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 +43,7 @@ tokenizer = AutoTokenizer.from_pretrained(
 )
 
 # Load the fine-tuned model
-ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir).to("cuda")  # Move model to CUDA
+ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800", cache_dir=cache_dir)
 
 def formatting_func(job_description):
     text = f"### The job description: {job_description}\n ### The skills: "
@@ -52,7 +52,7 @@ def formatting_func(job_description):
 @spaces.GPU  # Decorate the function to ensure it uses GPU
 def generate_skills(job_description):
     formatted_text = formatting_func(job_description)
-    model_input = tokenizer(formatted_text, return_tensors="pt").to("cuda")  # Use CUDA for GPU support
+    model_input = tokenizer(formatted_text, return_tensors="pt").to("cuda")  # Ensure input is on CUDA
 
     ft_model.eval()
     with torch.no_grad():