Create app.py

app.py

@@ -0,0 +1,378 @@
+import gradio as gr
+# --- Import necessary classes ---
+from transformers import (
+    TrOCRProcessor,
+    VisionEncoderDecoderModel,
+    pipeline,
+    AutoTokenizer,
+    RobertaForSequenceClassification,
+    AutoConfig  # <--- Import AutoConfig
+)
+# ---
+from PIL import Image
+import traceback
+import warnings
+import json
+import os
+import shutil # Not used directly, but keep for potential manual use
+
+# --- Model IDs ---
+TROCR_MODELS = {
+    "Printed Text": "microsoft/trocr-large-printed",
+    "Handwritten":  "microsoft/trocr-large-handwritten",
+}
+DETECTOR_MODEL_ID = "SuperAnnotate/roberta-large-llm-content-detector"
+print(f"Using AI Detector Model: {DETECTOR_MODEL_ID}")
+
+# --- Load OCR Models (no changes here) ---
+print("Loading OCR models...")
+OCR_PIPELINES = {}
+for name, model_id in TROCR_MODELS.items():
+    try:
+        proc = TrOCRProcessor.from_pretrained(model_id)
+        mdl  = VisionEncoderDecoderModel.from_pretrained(model_id)
+        OCR_PIPELINES[name] = (proc, mdl)
+        print(f"Loaded {name} OCR model.")
+    except Exception as e:
+        print(f"Error loading OCR model {name} ({model_id}): {e}")
+
+# --- Explicitly load config, tokenizer, and model ---
+print(f"Loading AI detector components ({DETECTOR_MODEL_ID})...")
+DETECTOR_PIPELINE = None
+detector_tokenizer = None
+detector_model = None
+try:
+    # 1. Load Configuration FIRST
+    print("Loading detector config...")
+    detector_config = AutoConfig.from_pretrained(DETECTOR_MODEL_ID)
+    print(f"Loaded config. Expected hidden size: {detector_config.hidden_size}") # Should be 1024
+
+    # Add an assertion to halt if config is wrong (optional but helpful)
+    if detector_config.hidden_size != 1024:
+         raise ValueError(f"Loaded config specifies hidden size {detector_config.hidden_size}, but expected 1024 for roberta-large. Check cache for {DETECTOR_MODEL_ID}.")
+
+    # 2. Load Tokenizer
+    print("Loading detector tokenizer...")
+    detector_tokenizer = AutoTokenizer.from_pretrained(DETECTOR_MODEL_ID)
+
+    # 3. Load Model using the specific class AND the loaded config
+    print("Loading detector model with loaded config...")
+    detector_model = RobertaForSequenceClassification.from_pretrained(
+        DETECTOR_MODEL_ID,
+        config=detector_config # <--- Pass the loaded config
+    )
+    print("AI detector model and tokenizer loaded successfully.")
+
+    # 4. Create Pipeline
+    print("Creating AI detector pipeline...")
+    DETECTOR_PIPELINE = pipeline(
+        "text-classification",
+        model=detector_model,
+        tokenizer=detector_tokenizer,
+        top_k=None
+    )
+    print("Created AI detector pipeline.")
+
+    # --- Optional: Label Test (keep from previous version) ---
+    if DETECTOR_PIPELINE:
+        try:
+            print("Testing detector pipeline labels...")
+            sample_output = DETECTOR_PIPELINE("This is a reasonably long test sentence to check the model labels.", truncation=True)
+            print(f"Sample detector output structure: {sample_output}")
+            # ... (rest of label testing code) ...
+            if sample_output and isinstance(sample_output, list) and len(sample_output) > 0:
+                 if isinstance(sample_output[0], list) and len(sample_output[0]) > 0:
+                     labels = [item.get('label', 'N/A') for item in sample_output[0] if isinstance(item, dict)]
+                     print(f"Detected labels from sample run: {labels}")
+                 elif isinstance(sample_output[0], dict):
+                     labels = [item.get('label', 'N/A') for item in sample_output if isinstance(item, dict)]
+                     print(f"Detected labels from sample run (non-nested): {labels}")
+                 if detector_model and detector_model.config and detector_model.config.id2label:
+                     print(f"Labels from model config: {detector_model.config.id2label}") # Should show {0: 'Human', 1: 'AI'}
+        except Exception as test_e:
+            print(f"Could not perform detector label test: {test_e}")
+            traceback.print_exc()
+
+
+except Exception as e:
+    print(f"CRITICAL Error loading AI detector components ({DETECTOR_MODEL_ID}): {e}")
+    traceback.print_exc()
+
+    # --- Simplified Cache Clearing Suggestion ---
+    # Get cache path using environment variable or default
+    hf_home = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface"))
+    hub_cache_path = os.path.join(hf_home, "hub") # Models are usually in the 'hub' subfolder
+
+    print("\n--- TROUBLESHOOTING SUGGESTION ---")
+    print(f"The model loading failed: {e}")
+    print("\nThis *strongly* indicates a problem with the cached files for this model.")
+    print("The most likely solution is to MANUALLY clear the cache for this model.")
+    print(f"\n1. Stop this application.")
+    print(f"2. Go to your Hugging Face hub cache directory (usually found under '{hub_cache_path}').")
+    print(f"   (If you've set HF_HOME environment variable, check there instead: '{hf_home}')")
+    # Construct the model-specific cache folder name
+    model_cache_folder_name = f"models--{DETECTOR_MODEL_ID.replace('/', '--')}"
+    print(f"3. Delete the specific folder for this model: '{model_cache_folder_name}'")
+    print(f"   Full path example: {os.path.join(hub_cache_path, model_cache_folder_name)}")
+    print(f"4. Restart the application. This will force a fresh download.")
+    print("\nMake sure no other applications are using the cache while deleting.")
+    print("--- END TROUBLESHOOTING ---")
+    # ---
+
+    # DETECTOR_PIPELINE remains None
+
+# --- Functions get_ai_and_human_scores, analyze_image, classify_text remain the same ---
+# (Ensure get_ai_and_human_scores correctly handles "AI" and "Human" based on config)
+def get_ai_and_human_scores(results):
+    """
+    Processes detector results to get likelihood scores for both AI and Human classes.
+    Handles various label formats including 'AI'/'Human', 'LABEL_0'/'LABEL_1', etc.
+    Returns:
+        tuple: (ai_display_string, human_display_string)
+    """
+    ai_prob = 0.0
+    human_prob = 0.0
+    status_message = "Status: Initializing..." # Default status
+
+    if not results:
+        print("Warning: Received empty results for AI detection.")
+        status_message = "Error: No results received"
+        return status_message, "N/A"
+
+    # Handle potential nested list structure
+    score_list = []
+    if isinstance(results, list) and len(results) > 0:
+        if isinstance(results[0], list) and len(results[0]) > 0:
+            score_list = results[0]
+        elif isinstance(results[0], dict):
+            score_list = results
+        else:
+            status_message = f"Error: Unexpected detector output format (inner list type: {type(results[0])})"
+            print(f"Warning: {status_message}. Results[0]: {results[0]}")
+            return status_message, "N/A"
+    else:
+        status_message = f"Error: Unexpected detector output format (outer type: {type(results)})"
+        print(f"Warning: {status_message}. Results: {results}")
+        return status_message, "N/A"
+
+    # Build label→score map (uppercase labels for robust matching)
+    lbl2score = {}
+    parse_errors = []
+    for entry in score_list:
+        if isinstance(entry, dict) and "label" in entry and "score" in entry:
+            try:
+                score = float(entry["score"])
+                lbl2score[entry["label"].upper()] = score
+            except (ValueError, TypeError):
+                parse_errors.append(f"Invalid score format: {entry}")
+        else:
+            parse_errors.append(f"Invalid entry format: {entry}")
+
+    if parse_errors:
+        print(f"Warning: Encountered parsing errors in score list: {parse_errors}")
+
+    if not lbl2score:
+        status_message = "Error: Could not parse any valid scores from detector output"
+        print(f"Warning: {status_message}. Score list was: {score_list}")
+        return status_message, "N/A"
+
+    label_keys_found = ", ".join(lbl2score.keys())
+    found_pair = False
+    inferred = False
+
+    # --- Determine AI and Human probabilities based on labels ---
+    upper_keys = lbl2score.keys()
+
+    # Prioritize AI/HUMAN as per model config
+    if "AI" in upper_keys and "HUMAN" in upper_keys:
+        ai_prob = lbl2score["AI"]
+        human_prob = lbl2score["HUMAN"]
+        found_pair = True
+        status_message = "OK (Used AI/HUMAN labels)"
+    # Fallbacks
+    elif "LABEL_1" in upper_keys and "LABEL_0" in upper_keys:
+        ai_prob = lbl2score["LABEL_1"]
+        human_prob = lbl2score["LABEL_0"]
+        found_pair = True
+        status_message = "OK (Used LABEL_1/LABEL_0 - Check Mapping)"
+        print("Warning: Used fallback LABEL_1/LABEL_0. Config expects AI/HUMAN.")
+    # Add other fallbacks if necessary (FAKE/REAL, MACHINE/HUMAN)
+
+    # Inference logic
+    if not found_pair:
+        if "AI" in upper_keys:
+            ai_prob = lbl2score["AI"]
+            human_prob = max(0.0, 1.0 - ai_prob)
+            inferred = True
+            status_message = "OK (Inferred from AI label)"
+        elif "HUMAN" in upper_keys:
+            human_prob = lbl2score["HUMAN"]
+            ai_prob = max(0.0, 1.0 - human_prob)
+            inferred = True
+            status_message = "OK (Inferred from HUMAN label)"
+        # Add fallback inference if needed
+
+        if not inferred:
+             status_message = f"Error: Could not determine AI/Human pair from labels [{label_keys_found}]"
+             print(f"Warning: {status_message}")
+
+    # --- Format output strings ---
+    ai_display_str = f"{ai_prob*100:.2f}%"
+    human_display_str = f"{human_prob*100:.2f}%"
+
+    if "Error:" in status_message:
+         ai_display_str = status_message
+         human_display_str = "N/A"
+
+    print(f"Score Status: {status_message}. AI={ai_display_str}, Human={human_display_str}")
+    return ai_display_str, human_display_str
+
+# --- analyze_image function (no changes needed) ---
+def analyze_image(image: Image.Image, ocr_choice: str):
+    """Performs OCR and AI Content Detection, returns both AI and Human %."""
+    extracted = ""
+    ai_result_str = "N/A"
+    human_result_str = "N/A"
+    status_update = "Awaiting input..."
+
+    if image is None:
+        status_update = "Please upload an image first."
+        return extracted, ai_result_str, human_result_str, status_update
+    if not ocr_choice or ocr_choice not in TROCR_MODELS:
+        status_update = "Please select a valid OCR model."
+        return extracted, ai_result_str, human_result_str, status_update
+    if OCR_PIPELINES.get(ocr_choice) is None:
+         return "", "N/A", "N/A", f"Error: OCR model '{ocr_choice}' failed to load or is unavailable."
+    if DETECTOR_PIPELINE is None:
+        return "", "N/A", "N/A", f"Critical Error: AI Detector model ({DETECTOR_MODEL_ID}) failed during startup. Check logs for details (possible cache issue?)."
+
+    try:
+        status_update = f"Processing with {ocr_choice} OCR..."
+        print(status_update)
+        proc, mdl = OCR_PIPELINES[ocr_choice]
+        if image.mode != "RGB": image = image.convert("RGB")
+        pix = proc(images=image, return_tensors="pt").pixel_values
+        tokens = mdl.generate(pix, max_length=1024)
+        extracted = proc.batch_decode(tokens, skip_special_tokens=True)[0]
+        extracted = extracted.strip()
+
+        if not extracted:
+             status_update = "OCR completed, but no text was extracted."
+             print(status_update)
+             return extracted, "N/A", "N/A", status_update
+
+        status_update = f"Detecting AI/Human content in {len(extracted)} characters..."
+        print(status_update)
+        results = DETECTOR_PIPELINE(extracted)
+
+        ai_result_str, human_result_str = get_ai_and_human_scores(results)
+
+        if "Error:" in ai_result_str:
+             status_update = ai_result_str
+        else:
+             status_update = "Analysis complete."
+        print(f"Final Status: {status_update}")
+
+        return extracted, ai_result_str, human_result_str, status_update
+
+    except Exception as e:
+        error_msg = f"Error during image analysis: {e}"
+        print(error_msg)
+        traceback.print_exc()
+        status_update = error_msg
+        return extracted, "Error", "Error", status_update
+
+
+# --- classify_text function (no changes needed) ---
+def classify_text(text: str):
+    """Classifies provided text, returning both AI and Human %."""
+    ai_result_str = "N/A"
+    human_result_str = "N/A"
+
+    if DETECTOR_PIPELINE is None:
+         return f"Critical Error: AI Detector model ({DETECTOR_MODEL_ID}) failed during startup. Check logs for details (possible cache issue?).", "N/A"
+    if not text or text.isspace():
+        return "Please enter some text.", "N/A"
+
+    print("Classifying text...")
+    try:
+        results = DETECTOR_PIPELINE(text)
+
+        ai_result_str, human_result_str = get_ai_and_human_scores(results)
+
+        if "Error:" not in ai_result_str:
+            print("Classification complete.")
+        else:
+            print(f"Classification completed with issues: {ai_result_str}")
+
+        return ai_result_str, human_result_str
+
+    except Exception as e:
+        error_msg = f"Error during text classification: {e}"
+        print(error_msg)
+        traceback.print_exc()
+        return error_msg, "Error"
+
+
+# --- Gradio Interface (no changes needed) ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        f"""
+        ## OCR + AI/Human Content Detection
+        Upload an image or paste text. The tool extracts text via OCR (if image) and analyzes it
+        using an AI content detector (`{DETECTOR_MODEL_ID}`)
+        to estimate the likelihood of it being AI-generated vs. Human-written.
+        **Disclaimer:** AI content detection is challenging and not 100% accurate. These likelihoods
+        are estimates based on the model's training data and may not be definitive.
+        Performance varies with text type, length, and AI generation methods.
+        **Label Assumption:** Uses the model's configured labels (`AI`/`Human`). Fallbacks for other label formats are included but may be less reliable if the model deviates from its configuration.
+        """
+    )
+    with gr.Tab("Analyze Image"):
+        with gr.Row():
+            with gr.Column(scale=2):
+                img_in = gr.Image(type="pil", label="Upload Image", sources=["upload", "clipboard"])
+            with gr.Column(scale=1):
+                ocr_dd = gr.Dropdown(
+                    list(TROCR_MODELS.keys()), label="1. Select OCR Model", info="Choose based on text type in image."
+                )
+                run_btn = gr.Button("2. Analyze Image", variant="primary")
+                status_img = gr.Label(value="Awaiting image analysis...", label="Status")
+
+        with gr.Row():
+            text_out_img = gr.Textbox(label="Extracted Text", lines=10, interactive=False)
+            with gr.Column(scale=1):
+                 ai_out_img = gr.Textbox(label="AI Likelihood %", interactive=False)
+            with gr.Column(scale=1):
+                 human_out_img = gr.Textbox(label="Human Likelihood %", interactive=False)
+
+        run_btn.click(
+            fn=analyze_image,
+            inputs=[img_in, ocr_dd],
+            outputs=[text_out_img, ai_out_img, human_out_img, status_img],
+            queue=True
+        )
+
+    with gr.Tab("Classify Text"):
+         with gr.Column():
+            text_in_classify = gr.Textbox(label="Paste or type text here", lines=10)
+            classify_btn = gr.Button("Classify Text", variant="primary")
+            with gr.Row():
+                 with gr.Column(scale=1):
+                    ai_out_classify = gr.Textbox(label="AI Likelihood %", interactive=False)
+                 with gr.Column(scale=1):
+                    human_out_classify = gr.Textbox(label="Human Likelihood %", interactive=False)
+
+            classify_btn.click(
+                fn=classify_text,
+                inputs=[text_in_classify],
+                outputs=[ai_out_classify, human_out_classify],
+                queue=True
+            )
+
+    gr.HTML(f"<footer style='text-align:center; margin-top: 20px; color: grey;'>Powered by TrOCR & {DETECTOR_MODEL_ID}</footer>")
+
+
+if __name__ == "__main__":
+    print("Starting Gradio demo...")
+    demo.launch(share=False, server_name="0.0.0.0")