Update app.py

app.py

@@ -12,9 +12,10 @@ from accelerate import Accelerator
 import bitsandbytes
 import sentencepiece
 import huggingface_hub
+from transformers import TrainingArguments, Trainer
 
 # Retrieve HF_TOKEN from Hugging Face Space secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
+HF_TOKEN = os.getenv("HF_TOKEN:levi put token here")  # Token expected as env variable 'HF_TOKEN'
 if not HF_TOKEN:
     raise ValueError("HF_TOKEN not found in environment variables. Please set it in Hugging Face Space secrets under 'Settings' > 'Secrets'.")
 
@@ -25,11 +26,6 @@ huggingface_hub.login(token=HF_TOKEN)
 MODEL_ID = "meta-llama/Llama-2-7b-hf"
 tokenizer = LlamaTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Add padding token if it doesn't exist
-if tokenizer.pad_token is None:
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    model.resize_token_embeddings(len(tokenizer))
-
 # Check CUDA and enable Flash Attention if supported
 use_flash_attention = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8
 model = LlamaForCausalLM.from_pretrained(
@@ -40,6 +36,11 @@ model = LlamaForCausalLM.from_pretrained(
     load_in_8bit=True
 )
 
+# Add padding token if it doesn't exist and resize embeddings
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+    model.resize_token_embeddings(len(tokenizer))
+
 # Prepare model for LoRA training
 model = prepare_model_for_kbit_training(model)
 peft_config = LoraConfig(
@@ -58,6 +59,7 @@ def train_ui(files):
     try:
         # Process multiple PDFs or JSON
         raw_text = ""
+        dataset = None  # Initialize dataset as None
         for file in files:
             if file.name.endswith(".pdf"):
                 with pdfplumber.open(file.name) as pdf:
@@ -152,4 +154,5 @@ with gr.Blocks(title="Healthcare Fraud Detection Fine-Tuning") as demo:
     output = gr.Textbox(label="Training Status", lines=5)
     train_button.click(fn=train_ui, inputs=file_input, outputs=output)
 
+# Launch the Gradio app
 demo.launch()