Update app.py

app.py

@@ -1,134 +1,133 @@
-# app.py (corrected version)
-
-# Handle missing dependencies first
-try:
-    import gradio as gr
-    from transformers import LlamaForCausalLM, LlamaTokenizer, Trainer, TrainingArguments
-    import datasets
-    import torch
-    import json
-    import os
-    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-    from accelerate import Accelerator
-    import bitsandbytes
-    import sentencepiece  # Added for Llama tokenizer
-except ImportError as e:
-    missing_package = str(e).split("'")[-2]  # Extract the missing package name
-    if "accelerate" in missing_package:
-        os.system(f'pip install "accelerate>=0.26.0"')
-    elif "sentencepiece" in missing_package:
-        os.system(f'pip install "sentencepiece"')
-    else:
-        os.system(f'pip install "{missing_package}"')
-    # Re-import after installation
-    import gradio as gr
-    from transformers import LlamaForCausalLM, LlamaTokenizer, Trainer, TrainingArguments
-    import datasets
-    import torch
-    import json
-    import os
-    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
-    from accelerate import Accelerator
-    import bitsandbytes
-    import sentencepiece
+# app.py
+
+import gradio as gr
+from transformers import LlamaForCausalLM, LlamaTokenizer
+import datasets
+import torch
+import json
+import os
+import pdfplumber
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from accelerate import Accelerator
+import bitsandbytes
+import sentencepiece
+import huggingface_hub
+
+# Retrieve HF_TOKEN from Hugging Face Space secrets
+HF_TOKEN = os.getenv("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'.")
+
+# Authenticate with Hugging Face
+huggingface_hub.login(token=HF_TOKEN)
 
 # Model setup
-MODEL_ID = "meta-llama/Llama-2-7b-hf"  # Use Llama-2-7b; switch to "meta-llama/Llama-3-8b-hf" for Llama 3
-tokenizer = LlamaTokenizer.from_pretrained(MODEL_ID)
+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 (required for Llama models)
+# 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 if CUDA is available to enable Flash Attention 2
-use_flash_attention = torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 8  # Ampere or newer (e.g., A100)
-
-# Load the model with optimizations for Llama
+# 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(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,  # Better for A100 GPUs, falls back to float16 on CPU
+    torch_dtype=torch.bfloat16,
     device_map="auto",
-    use_flash_attention_2=use_flash_attention,  # Only enable if GPU supports it
-    load_in_8bit=True  # Quantization for memory efficiency
+    use_flash_attention_2=use_flash_attention,
+    load_in_8bit=True
 )
 
-# Prepare the model for training with LoRA (more memory-efficient)
+# Prepare model for LoRA training
 model = prepare_model_for_kbit_training(model)
-
-# LoRA configuration
 peft_config = LoraConfig(
-    r=16,               # Rank
-    lora_alpha=32,      # Alpha
-    lora_dropout=0.05,  # Dropout
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
     bias="none",
     task_type="CAUSAL_LM",
-    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]  # Attention modules for Llama
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj"]
 )
-
 model = get_peft_model(model, peft_config)
-model.print_trainable_parameters()  # Print percentage of trainable parameters
+model.print_trainable_parameters()
 
-# Function to process uploaded JSON and train
-def train_ui_tars(file):
+# Function to process uploaded files and train
+def train_ui(files):
     try:
-        # Step 1: Load and preprocess the uploaded JSON file
-        with open(file.name, "r", encoding="utf-8") as f:
-            raw_data = json.load(f)
-       
-        # Extract training pairs or use flat structure
-        training_data = raw_data.get("training_pairs", raw_data)
-       
-        # Save fixed JSON to avoid issues
-        fixed_json_path = "fixed_fraud_data.json"
-        with open(fixed_json_path, "w", encoding="utf-8") as f:
-            json.dump(training_data, f, indent=4)
-       
-        # Load dataset
-        dataset = datasets.load_dataset("json", data_files=fixed_json_path)
-       
-        # Step 2: Tokenize dataset with Llama-compatible context length
+        # Process multiple PDFs or JSON
+        raw_text = ""
+        for file in files:
+            if file.name.endswith(".pdf"):
+                with pdfplumber.open(file.name) as pdf:
+                    for page in pdf.pages:
+                        raw_text += page.extract_text() or ""
+            elif file.name.endswith(".json"):
+                with open(file.name, "r", encoding="utf-8") as f:
+                    raw_data = json.load(f)
+                    training_data = raw_data.get("training_pairs", raw_data)
+                    with open("temp_fraud_data.json", "w", encoding="utf-8") as f:
+                        json.dump({"training_pairs": training_data}, f)
+                    dataset = datasets.load_dataset("json", data_files="temp_fraud_data.json")
+
+        if not raw_text and not dataset:
+            return "Error: No valid PDF or JSON data found."
+
+        # Create training pairs from PDFs if no JSON
+        if raw_text:
+            def create_training_pairs(text):
+                pairs = []
+                if "Haloperidol" in text and "daily" in text.lower():
+                    pairs.append({
+                        "input": "Patient received Haloperidol daily. Is this overmedication?",
+                        "output": "Yes, daily Haloperidol use without documented severe psychosis or failed alternatives may indicate overmedication, violating CMS guidelines."
+                    })
+                if "Lorazepam" in text and "frequent" in text.lower():
+                    pairs.append({
+                        "input": "Care logs show frequent Lorazepam use with a 90-day supply. Is this suspicious?",
+                        "output": "Yes, frequent use with a large supply suggests potential overuse or mismanagement, a fraud indicator."
+                    })
+                return pairs
+            training_data = create_training_pairs(raw_text)
+            with open("temp_fraud_data.json", "w") as f:
+                json.dump({"training_pairs": training_data}, f)
+            dataset = datasets.load_dataset("json", data_files="temp_fraud_data.json")
+
+        # Tokenization function
         def tokenize_data(example):
-            # Format input for Llama (instruction-following style)
             formatted_text = f"<s>[INST] {example['input']} [/INST] {example['output']}</s>"
-            inputs = tokenizer(
-                formatted_text,
-                padding="max_length",
-                truncation=True,
-                max_length=2048,  # Llama 2 context length; adjust to 8192 for Llama 3 if needed
-                return_tensors="pt"
-            )
+            inputs = tokenizer(formatted_text, padding="max_length", truncation=True, max_length=4096, return_tensors="pt")
             inputs["labels"] = inputs["input_ids"].clone()
             return {k: v.squeeze(0) for k, v in inputs.items()}
-       
+
         tokenized_dataset = dataset["train"].map(tokenize_data, batched=True, remove_columns=dataset["train"].column_names)
-       
-        # Step 3: Training setup
+
+        # Training setup
         training_args = TrainingArguments(
-            output_dir="./fine_tuned_llama",
-            per_device_train_batch_size=4,  # Increased for better efficiency
-            gradient_accumulation_steps=8,  # To handle larger effective batch size
-            evaluation_strategy="no",
+            output_dir="./fine_tuned_llama_healthcare",
+            per_device_train_batch_size=4,
+            gradient_accumulation_steps=8,
+            eval_strategy="no",
             save_strategy="epoch",
             save_total_limit=2,
-            num_train_epochs=3,
+            num_train_epochs=5,
             learning_rate=2e-5,
             weight_decay=0.01,
             logging_dir="./logs",
             logging_steps=10,
-            bf16=True,  # Use bfloat16 for A100 GPUs, falls back to float16 on CPU
-            gradient_checkpointing=True,  # Memory optimization
+            bf16=True,
+            gradient_checkpointing=True,
             optim="adamw_torch",
             warmup_steps=100,
         )
-       
-        # Custom data collator for Llama
+
         def custom_data_collator(features):
-            batch = {
+            return {
                 "input_ids": torch.stack([f["input_ids"] for f in features]),
                 "attention_mask": torch.stack([f["attention_mask"] for f in features]),
                 "labels": torch.stack([f["labels"] for f in features]),
             }
-            return batch
 
         trainer = Trainer(
             model=model,
@@ -136,28 +135,21 @@ def train_ui_tars(file):
             train_dataset=tokenized_dataset,
             data_collator=custom_data_collator,
         )
-       
-        # Step 4: Start training
         trainer.train()
-       
-        # Step 5: Save the model
-        model.save_pretrained("./fine_tuned_llama")
-        tokenizer.save_pretrained("./fine_tuned_llama")
-       
-        return "Training completed successfully! Model saved to ./fine_tuned_llama"
-   
+        model.save_pretrained("./fine_tuned_llama_healthcare")
+        tokenizer.save_pretrained("./fine_tuned_llama_healthcare")
+        return "Training completed! Model saved to ./fine_tuned_llama_healthcare"
+
     except Exception as e:
-        return f"Error: {str(e)}"
+        return f"Error: {str(e)}. Please check file format, dependencies, or HF_TOKEN."
 
 # Gradio UI
-with gr.Blocks(title="Model Fine-Tuning Interface") as demo:
-    gr.Markdown("# Llama Fraud Detection Fine-Tuning UI")
-    gr.Markdown("Upload a JSON file with 'input' and 'output' pairs to fine-tune the Llama model on your fraud dataset.")
-   
-    file_input = gr.File(label="Upload Fraud Dataset (JSON)")
+with gr.Blocks(title="Healthcare Fraud Detection Fine-Tuning") as demo:
+    gr.Markdown("# Fine-Tune LLaMA 2 for Healthcare Fraud Analysis")
+    gr.Markdown("Upload PDFs (e.g., care logs, medication records) or a JSON file with training pairs.")
+    file_input = gr.File(label="Upload Files (PDF/JSON)", file_count="multiple")
     train_button = gr.Button("Start Fine-Tuning")
-    output = gr.Textbox(label="Training Status")
-   
-    train_button.click(fn=train_ui_tars, inputs=file_input, outputs=output)
+    output = gr.Textbox(label="Training Status", lines=5)
+    train_button.click(fn=train_ui, inputs=file_input, outputs=output)
 
 demo.launch()