Update app.py

app.py

@@ -1,8 +1,8 @@
 # updated_app.py
-# Enhanced Gradio app for Llama 4 Maverick healthcare fraud detection
+# Enhanced Gradio app for Llama 4 Maverick healthcare fraud detection (text-only)
 
 import gradio as gr
-from transformers import AutoProcessor, Llama4ForConditionalGeneration
+from transformers import AutoTokenizer, Llama4ForConditionalGeneration
 import datasets
 import torch
 import json
@@ -39,9 +39,13 @@ huggingface_hub.login(token=LLama)
 
 # Model setup
 MODEL_ID = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
-processor = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# Load model with FP8 quantization to fit in 80 GB VRAM
+# Add padding token if it doesn't exist
+if tokenizer.pad_token is None:
+    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+
+# Load model with 8-bit quantization to fit in 80 GB VRAM
 model = Llama4ForConditionalGeneration.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
@@ -157,18 +161,8 @@ def train_ui(files):
             dataset = datasets.load_dataset("json", data_files="temp_fraud_data.json")
 
         def tokenize_data(example):
-            messages = [
-                {
-                    "role": "user",
-                    "content": [{"type": "text", "text": example['input']}]
-                },
-                {
-                    "role": "assistant",
-                    "content": [{"type": "text", "text": example['output']}]
-                }
-            ]
-            formatted_text = processor.apply_chat_template(messages, add_generation_prompt=False)
-            inputs = processor(formatted_text, padding="max_length", truncation=True, max_length=4096, return_tensors="pt")
+            formatted_text = f"<s>[INST] {example['input']} [/INST] {example['output']}</s>"
+            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()}
 
@@ -208,7 +202,7 @@ def train_ui(files):
 
         trainer.train()
         model.save_pretrained("./fine_tuned_llama4_healthcare")
-        processor.save_pretrained("./fine_tuned_llama4_healthcare")
+        tokenizer.save_pretrained("./fine_tuned_llama4_healthcare")
         return f"Training completed with {len(tokenized_dataset)} examples! Model saved to ./fine_tuned_llama4_healthcare"
 
     except Exception as e:
@@ -232,7 +226,7 @@ def analyze_document_ui(files):
         if not raw_text:
             return "Error: Could not extract text from the PDF. The file may be corrupt or contain only images."
         
-        analyzer = HealthcareFraudAnalyzer(model, processor)
+        analyzer = HealthcareFraudAnalyzer(model, tokenizer)
         results = analyzer.analyze_document(raw_text)
         return results["summary"]