token

app.py

@@ -1,4 +1,4 @@
-import os 
+import os  # Ensure os is imported
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import gradio as gr
@@ -21,19 +21,29 @@ 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'.")
 
+# Define cache directory
+cache_dir = "./cache"
+
 # Load the base model with the updated 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
-)
+    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(base_model_id, add_bos_token=True, trust_remote_code=True, token=huggingface_token)
+tokenizer = AutoTokenizer.from_pretrained(
+    base_model_id, 
+    add_bos_token=True, 
+    trust_remote_code=True, 
+    token=huggingface_token,
+    cache_dir=cache_dir  # Specify cache directory
+)
 
 # Load the fine-tuned model
-ft_model = PeftModel.from_pretrained(base_model, "checkpoint-2800")
+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: "
@@ -65,4 +75,4 @@ inputs = gr.Textbox(lines=10, label="Job description:", placeholder="Enter or pa
 outputs = gr.Textbox(label="Required skills:", placeholder="The required skills will be displayed here...")
 
 gr.Interface(fn=generate_skills, inputs=inputs, outputs=outputs, title="Job Skills Analysis",
-             description="Paste the job description in the text box below and the model will show the required skills for candidates.").launch()
+             description="Paste the job description in the text box below and the model will show the required skills for candidates.").launch(share=True)  # Set share=True to create a public link