Update app.py

app.py

@@ -1,5 +1,5 @@
 # app.py
-# Enhanced Gradio app for Llama 4 Maverick healthcare fraud detection (text-only with CPU offloading)
+# Gradio app for Llama 4 Maverick healthcare fraud detection (text-only with CPU offloading)
 
 import gradio as gr
 from transformers import AutoTokenizer, Llama4ForConditionalGeneration
@@ -23,15 +23,15 @@ except LookupError:
 # Import the HealthcareFraudAnalyzer
 from document_analyzer import HealthcareFraudAnalyzer
 
-# Debug: Print environment variables to verify 'LLama' is present
+# Debug: Print environment variables
 print("Environment variables:", dict(os.environ))
 
-# Retrieve the token from Hugging Face Space secrets
+# Retrieve the token from secrets
 LLama = os.getenv("LLama")
 if not LLama:
     raise ValueError("LLama token not found. Set it in Hugging Face Space secrets as 'LLama'.")
 
-# Debug: Print token (first 5 chars for security, remove in production)
+# Debug: Print token (first 5 chars)
 print(f"Retrieved LLama token: {LLama[:5]}...")
 
 # Authenticate with Hugging Face
@@ -41,20 +41,19 @@ huggingface_hub.login(token=LLama)
 MODEL_ID = "meta-llama/Llama-4-Maverick-17B-128E-Instruct"
 tokenizer = AutoTokenizer.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]'})
 
-# Custom device map to offload some layers to CPU
+# Device map for CPU offloading
 device_map = {
     "model.embed_tokens": 0,
-    "model.layers.0-15": 0,  # Keep first 16 layers on GPU
-    "model.layers.16-31": "cpu",  # Offload remaining layers to CPU
+    "model.layers.0-15": 0,
+    "model.layers.16-31": "cpu",
     "model.norm": 0,
     "lm_head": 0
 }
 
-# Load model with 8-bit quantization and CPU offloading
+# Load model with 8-bit quantization
 model = Llama4ForConditionalGeneration.from_pretrained(
     MODEL_ID,
     torch_dtype=torch.bfloat16,
@@ -64,7 +63,7 @@ model = Llama4ForConditionalGeneration.from_pretrained(
     attn_implementation="flex_attention"
 )
 
-# Prepare model for LoRA training
+# Prepare for LoRA training
 model = prepare_model_for_kbit_training(model)
 peft_config = LoraConfig(
     r=16,
@@ -77,73 +76,53 @@ peft_config = LoraConfig(
 model = get_peft_model(model, peft_config)
 model.print_trainable_parameters()
 
-# Function to create training pairs from document text
+# Function to create training pairs
 def extract_training_pairs_from_text(text):
     pairs = []
     patterns = [
-        # Medication patterns
-        (
-            r"(?i).*?\b(haloperidol|lorazepam|ativan)\b.*?\b(daily|routine|regular)\b.*?",
-            "Patient receives {} on a {} basis. Is this appropriate medication management?",
-            "This may indicate inappropriate medication management. Regular use of psychotropic medications without documented need assessment, behavior monitoring, and attempted dose reductions may violate care standards."
-        ),
-        # Documentation patterns
-        (
-            r"(?i).*?\b(missing|omitted|absent|lacking)\b.*?\b(documentation|records|logs|notes)\b.*?",
-            "Facility has {} {} for patient care. Is this a documentation concern?",
-            "Yes, incomplete documentation is a significant red flag. Missing records may indicate attempts to conceal care issues or fraudulent billing for services not provided."
-        ),
-        # Visitation patterns
-        (
-            r"(?i).*?\b(restrict|limit|prevent|block)\b.*?\b(visits|visitation|access|family)\b.*?",
-            "Facility {} family {} without documented medical necessity. Is this suspicious?",
-            "Yes, unjustified visitation restrictions may indicate attempts to conceal care issues and prevent family oversight. This can constitute fraud when facilities bill for care while violating resident rights."
-        ),
-        # Hospice patterns
-        (
-            r"(?i).*?\b(hospice|terminal|end.of.life)\b.*?\b(not|without|lacking)\b.*?\b(evidence|decline|documentation)\b.*?",
-            "Patient placed on {} care {} supporting {}. Is this fraudulent?",
-            "Yes, hospice enrollment without documented terminal decline may indicate Medicare fraud. Hospice certification requires genuine clinical determination of terminal status with prognosis of six months or less."
-        ),
-        # Contradictory documentation
-        (
-            r"(?i).*?\b(different|contradicts|conflicts|inconsistent)\b.*?\b(records|documentation|testimony|statements)\b.*?",
-            "Records show {} {} about patient condition. Is this fraudulent documentation?",
-            "Yes, contradictory documentation is a strong indicator of fraudulent record-keeping designed to misrepresent care quality or patient condition, particularly when official records differ from internal communications."
-        )
+        (r"(?i).*?\b(haloperidol|lorazepam|ativan)\b.*?\b(daily|routine|regular)\b.*?",
+         "Patient receives {} on a {} basis. Is this appropriate?",
+         "This may indicate inappropriate use. Regular psychotropic use without need assessment may violate standards."),
+        (r"(?i).*?\b(missing|omitted|absent|lacking)\b.*?\b(documentation|records|logs|notes)\b.*?",
+         "Facility has {} {} for care. Is this a concern?",
+         "Yes, incomplete records may indicate fraud or attempts to hide issues."),
+        (r"(?i).*?\b(restrict|limit|prevent|block)\b.*?\b(visits|visitation|access|family)\b.*?",
+         "Facility {} family {} without necessity. Is this suspicious?",
+         "Yes, restrictions may hide issues and constitute fraud when billing for care."),
+        (r"(?i).*?\b(hospice|terminal|end.of.life)\b.*?\b(not|without|lacking)\b.*?\b(evidence|decline|documentation)\b.*?",
+         "Patient on {} care {} supporting {}. Is this fraudulent?",
+         "Yes, hospice without documented decline may indicate Medicare fraud."),
+        (r"(?i).*?\b(different|contradicts|conflicts|inconsistent)\b.*?\b(records|documentation|testimony|statements)\b.*?",
+         "Records show {} {} about condition. Is this fraudulent?",
+         "Yes, contradictory records suggest fraudulent misrepresentation.")
     ]
 
     for pattern, input_template, output_text in patterns:
-        matches = re.finditer(pattern, text)
-        for match in matches:
+        for match in re.finditer(pattern, text):
             groups = match.groups()
             if len(groups) >= 2:
-                input_text = input_template.format(*groups)
-                pairs.append({
-                    "input": input_text,
-                    "output": output_text
-                })
+                pairs.append({"input": input_template.format(*groups), "output": output_text})
 
     if not pairs:
         if any(x in text.lower() for x in ["medication", "prescribed", "administered"]):
             pairs.append({
-                "input": "Medication records show inconsistencies in administration times. Is this concerning?",
-                "output": "Yes, inconsistent medication administration timing may indicate fraudulent documentation or medication mismanagement that could harm patients."
+                "input": "Medication records show inconsistent times. Is this concerning?",
+                "output": "Yes, inconsistent timing may indicate fraud or mismanagement."
             })
         if any(x in text.lower() for x in ["visit", "family", "spouse"]):
             pairs.append({
-                "input": "Staff documents family visits inconsistently. Is this suspicious?",
-                "output": "Yes, selective documentation of family visits indicates fraudulent record-keeping designed to create a false narrative about family involvement and patient responses."
+                "input": "Staff documents visits inconsistently. Is this suspicious?",
+                "output": "Yes, selective documentation suggests fraudulent record-keeping."
             })
         if any(x in text.lower() for x in ["hospice", "terminal", "prognosis"]):
             pairs.append({
-                "input": "Patient remained on hospice for extended period without documented decline. Is this Medicare fraud?",
-                "output": "Yes, maintaining hospice services without documented decline suggests fraudulent hospice certification to obtain Medicare benefits inappropriately."
+                "input": "Patient on hospice without decline. Is this fraud?",
+                "output": "Yes, lack of decline suggests fraudulent certification."
             })
 
     return pairs
 
-# Function to process uploaded files and train
+# Function to process files and train
 def train_ui(files):
     try:
         raw_text = ""
@@ -218,15 +197,15 @@ def train_ui(files):
     except Exception as e:
         return f"Error: {str(e)}. Please check file format, dependencies, or the LLama token."
 
-# Function to analyze uploaded document for fraud
+# Function to analyze documents
 def analyze_document_ui(files):
     try:
         if not files:
-            return "Error: No file uploaded. Please upload a PDF to analyze."
+            return "Error: No file uploaded. Please upload a PDF."
         
         file = files[0]
         if not file.name.endswith(".pdf"):
-            return "Error: Please upload a PDF file for analysis."
+            return "Error: Please upload a PDF file."
         
         raw_text = ""
         with pdfplumber.open(file.name) as pdf:
@@ -234,43 +213,42 @@ def analyze_document_ui(files):
                 raw_text += page.extract_text() or ""
         
         if not raw_text:
-            return "Error: Could not extract text from the PDF. The file may be corrupt or contain only images."
+            return "Error: Could not extract text from PDF."
         
         analyzer = HealthcareFraudAnalyzer(model, tokenizer)
         results = analyzer.analyze_document(raw_text)
         return results["summary"]
     
     except Exception as e:
-        return f"Error during document analysis: {str(e)}"
+        return f"Error during analysis: {str(e)}"
 
-# Gradio UI with training and analysis tabs
+# Gradio UI
 with gr.Blocks(title="Healthcare Fraud Detection Suite") as demo:
     gr.Markdown("# Healthcare Fraud Detection Suite")
     
     with gr.Tabs():
         with gr.TabItem("Fine-Tune Model"):
-            gr.Markdown("## Train Llama 4 for Healthcare Fraud Detection")
-            gr.Markdown("Upload PDFs (e.g., care logs, medication records) or a JSON file with training pairs.")
+            gr.Markdown("## Train Llama 4 for Fraud Detection")
+            gr.Markdown("Upload PDFs or JSON with training pairs.")
             train_file_input = gr.File(label="Upload Files (PDF/JSON)", file_count="multiple")
             train_button = gr.Button("Start Fine-Tuning")
             train_output = gr.Textbox(label="Training Status", lines=5)
             train_button.click(fn=train_ui, inputs=train_file_input, outputs=train_output)
         
         with gr.TabItem("Analyze Document"):
-            gr.Markdown("## Analyze Document for Healthcare Fraud Indicators")
-            gr.Markdown("Upload a PDF document to analyze for potential fraud, neglect, or abuse indicators.")
-            analyze_file_input = gr.File(label="Upload PDF Document")
+            gr.Markdown("## Analyze for Fraud Indicators")
+            gr.Markdown("Upload a PDF to scan for fraud, neglect, or abuse.")
+            analyze_file_input = gr.File(label="Upload PDF")
             analyze_button = gr.Button("Analyze Document")
             analyze_output = gr.Markdown(label="Analysis Results")
             analyze_button.click(fn=analyze_document_ui, inputs=analyze_file_input, outputs=analyze_output)
     
     gr.Markdown("""
     ### About This Tool
-    This tool uses Llama 4 Maverick to identify patterns of potential fraud, neglect, and abuse in healthcare documentation.
-    The fine-tuning tab allows model customization with your examples or automatic extraction from documents.
-    The analysis tab scans documents for suspicious patterns, generating detailed reports.
-    **Note:** All analysis is performed locally - no data is shared externally.
+    Uses Llama 4 Maverick to detect fraud in healthcare documents.
+    Fine-tune with custom data or analyze PDFs for suspicious patterns.
+    **Note:** All analysis is local - no data is shared.
     """)
 
-# Launch the Gradio app
+# Launch the app
 demo.launch()