app.py

import gradio as gr
import warnings
warnings.filterwarnings("ignore")

# Import utilities and configuration
from utils import BailOutcomePredictor, FairnessPredictor, DatasetLoader
from config import *

# Initialize predictors and dataset loader
print("Initializing models and dataset...")

bail_predictor = BailOutcomePredictor(
    bert_model_path=BAIL_BERT_MODEL_PATH,
    logreg_model_path=BAIL_LOGREG_MODEL_PATH,
    tfidf_path=BAIL_TFIDF_PATH,
    scaler_path=BAIL_SCALER_PATH
)

fairness_predictor = FairnessPredictor(
    bert_model_path=FAIRNESS_BERT_MODEL_PATH,
    logreg_model_path=FAIRNESS_LOGREG_MODEL_PATH,
    xgb_model_path=FAIRNESS_XGB_MODEL_PATH,
    preprocessor_paths={
        'fairness_tfidf_vectorizer': FAIRNESS_TFIDF_PATH,
        'fairness_feature_scaler': FAIRNESS_SCALER_PATH,
        'fairness_label_encoders': FAIRNESS_ENCODERS_PATH,
        'fairness_xgb_tfidf_vectorizer': FAIRNESS_XGB_TFIDF_PATH,
        'fairness_xgb_feature_scaler': FAIRNESS_XGB_SCALER_PATH,
        'fairness_xgb_label_encoders': FAIRNESS_XGB_ENCODERS_PATH
    }
)

dataset_loader = DatasetLoader(DATASET_PATH)

print("Initialization complete!")

# Prediction functions for Gradio interface
def predict_bail_bert(case_text):
    """Predict bail outcome using BERT model"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    return bail_predictor.predict_with_bert(case_text)

def predict_bail_logreg(case_text, crime_type, bail_type):
    """Predict bail outcome using Logistic Regression model"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    return bail_predictor.predict_with_logreg(case_text, crime_type, bail_type)

def predict_fairness_bert(case_text):
    """Predict fairness using BERT model"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    return fairness_predictor.predict_with_bert(case_text)

def predict_fairness_logreg(case_text, accused_gender):
    """Predict fairness using Logistic Regression model"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    return fairness_predictor.predict_with_logreg(case_text, accused_gender)

def predict_fairness_xgboost(case_text, accused_gender):
    """Predict fairness using XGBoost model"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    return fairness_predictor.predict_with_xgboost(case_text, accused_gender)

def compare_all_models(case_text, crime_type, accused_gender, bail_type):
    """Compare predictions from all models"""
    if not case_text.strip():
        return {"Error": "Please enter case details"}
    
    results = {
        "Bail Outcome": {
            "LegalBERT": predict_bail_bert(case_text),
            "Logistic Regression": predict_bail_logreg(case_text, crime_type, bail_type)
        },
        "Fairness Analysis": {
            "LegalBERT": predict_fairness_bert(case_text),
            "Logistic Regression": predict_fairness_logreg(case_text, accused_gender),
            "XGBoost": predict_fairness_xgboost(case_text, accused_gender)
        }
    }
    return results

def load_sample_case():
    """Load a sample case from a curated list of diverse cases"""
    # Curated sample cases representing different crime types, genders, and outcomes
    sample_cases = [
        {
            "case_text": "The accused was found in possession of 50 grams of heroin during a police raid. The investigating officer testified that the accused was caught red-handed with the contraband. The accused claims innocence and states that the drugs were planted. Medical reports show the accused has a history of substance abuse. The prosecution argues that the accused is likely to flee if granted bail due to the serious nature of the offense.",
            "gender": "Male",
            "crime": "Narcotics",
            "bail": "Regular",
            "region": "Delhi"
        },
        {
            "case_text": "The complainant alleges that the accused sexually assaulted her minor daughter. Medical examination confirms assault. The accused denies the charges and claims false implication due to property dispute with the complainant's family. The prosecution opposes bail citing the serious nature of the offense and possibility of the accused influencing witnesses in the locality.",
            "gender": "Male", 
            "crime": "Sexual Offense",
            "bail": "Regular",
            "region": "Maharashtra"
        },
        {
            "case_text": "The accused is alleged to have cheated investors of Rs. 2 crore through a fraudulent investment scheme. Multiple FIRs have been filed across different states. The accused argues that it was a genuine business venture that failed due to market conditions. The prosecution contends that the accused has properties worth crores and may flee the country if granted bail.",
            "gender": "Male",
            "crime": "Fraud or Cheating", 
            "bail": "Regular",
            "region": "Karnataka"
        },
        {
            "case_text": "The accused was involved in a minor road accident case. No serious injuries were reported. The accused has surrendered before the court and is cooperating with the investigation. The accused has no prior criminal history and has strong roots in the community. The prosecution does not oppose bail.",
            "gender": "Male",
            "crime": "Others",
            "bail": "Regular", 
            "region": "Tamil Nadu"
        },
        {
            "case_text": "The accused is alleged to have murdered her husband following domestic disputes. The prosecution claims it was a premeditated murder. The accused claims self-defense citing years of domestic violence. Neighbors testify about frequent fights between the couple. The accused has no prior criminal record and has elderly parents to take care of.",
            "gender": "Female",
            "crime": "Murder",
            "bail": "Regular",
            "region": "West Bengal"
        },
        {
            "case_text": "The accused is charged with harassing his daughter-in-law for dowry. The complainant alleges physical and mental torture for additional dowry demands. The accused denies the charges and claims the complaint is motivated by family property disputes. The complainant has submitted medical records showing injuries. The accused is aged and has no prior criminal history.",
            "gender": "Male",
            "crime": "Dowry Harassment",
            "bail": "Anticipatory",
            "region": "Uttar Pradesh"
        },
        {
            "case_text": "The accused allegedly robbed a jewelry shop at gunpoint and fled with gold worth Rs. 5 lakhs. CCTV footage shows a person matching the accused's description. The accused was arrested after a chase and some stolen items were recovered. The accused claims mistaken identity and has an alibi for the time of the incident.",
            "gender": "Male",
            "crime": "Theft or Robbery",
            "bail": "Regular",
            "region": "Gujarat" 
        },
        {
            "case_text": "The accused is alleged to have kidnapped a businessman's son for ransom. The victim was recovered safely after police investigation. The accused denies involvement and claims he was elsewhere during the incident. The prosecution argues that the accused is a flight risk due to the serious nature of the offense and his history of changing addresses frequently.",
            "gender": "Male", 
            "crime": "Kidnapping",
            "bail": "Regular",
            "region": "Haryana"
        }
    ]
    
    # Randomly select one of the curated cases
    import random
    selected_case = random.choice(sample_cases)
    
    return (
        selected_case["case_text"],
        selected_case["case_text"], 
        selected_case["case_text"],
        selected_case["gender"],
        selected_case["crime"],
        selected_case["bail"],
        selected_case["region"]
    )

# Create Gradio interface
with gr.Blocks(title="Indian Bail Judgment Analysis", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏛️ Indian Bail Judgment Analysis System
    
    **AI-powered analysis of Indian bail judgments** for bail outcome prediction and bias detection.
    
    ## 🚀 Quick Start Guide
    **New to the system?** Click the **🎲 Load Sample Case** button in any tab to explore with ready-made examples covering:
    - 8 different crime types (Narcotics, Sexual Offense, Fraud, Murder, etc.)
    - Various accused genders and demographics  
    - Different bail outcomes and regional variations
    - Real case scenarios from the dataset
    """)
    
    with gr.Tabs():
        # Bail Outcome Prediction Tab
        with gr.TabItem("🏛️ Bail Outcome Prediction"):
            gr.Markdown("### Predict whether bail will be granted or rejected")
            gr.Markdown("💡 **Tip**: Don't have a case ready? Click the '**Load Sample Case**' button below for instant examples with different crime types!")
            
            with gr.Row():
                with gr.Column():
                    case_input_bail = gr.Textbox(
                        label="Case Details",
                        placeholder="Enter case facts, legal issues, and other relevant details... OR click 'Load Sample Case' below!",
                        lines=10
                    )
                    
                    with gr.Row():
                        crime_type_bail = gr.Dropdown(
                            choices=CRIME_TYPES,
                            value="Unknown",
                            label="Crime Type"
                        )
                        bail_type_bail = gr.Dropdown(
                            choices=BAIL_TYPES,
                            value="Regular",
                            label="Bail Type"
                        )
                    
                    with gr.Row():
                        predict_bert_bail = gr.Button("Predict with LegalBERT", variant="primary")
                        predict_logreg_bail = gr.Button("Predict with LogReg", variant="secondary")
                        load_sample_bail = gr.Button("🎲 Load Sample Case", variant="primary", scale=1)
                
                with gr.Column():
                    result_bert_bail = gr.JSON(label="LegalBERT Results")
                    result_logreg_bail = gr.JSON(label="Logistic Regression Results")
        
        # Fairness Analysis Tab
        with gr.TabItem("⚖️ Fairness Analysis"):
            gr.Markdown("### Analyze potential bias in bail judgments")
            gr.Markdown("💡 **Tip**: Use the sample cases by clicking '**🎲 Load Sample Case**' to see how different models detect bias!")
            
            with gr.Row():
                with gr.Column():
                    case_input_fairness = gr.Textbox(
                        label="Case Details",
                        placeholder="Enter case facts, legal issues, and judgment reasoning... OR use sample cases below!",
                        lines=10
                    )
                    
                    with gr.Row():
                        accused_gender = gr.Dropdown(
                            choices=GENDERS,
                            value="Male",
                            label="Accused Gender"
                        )
                        region_fairness = gr.Dropdown(
                            choices=REGIONS,
                            value="Unknown",
                            label="Region"
                        )
                    
                    with gr.Row():
                        predict_bert_fairness = gr.Button("Predict with LegalBERT", variant="primary")
                        predict_logreg_fairness = gr.Button("Predict with LogReg", variant="secondary")
                        predict_xgb_fairness = gr.Button("Predict with XGBoost", variant="secondary")
                
                with gr.Column():
                    result_bert_fairness = gr.JSON(label="LegalBERT Fairness Results")
                    result_logreg_fairness = gr.JSON(label="Logistic Regression Fairness Results")
                    result_xgb_fairness = gr.JSON(label="XGBoost Fairness Results")
        
        # Model Comparison Tab
        with gr.TabItem("📊 Model Comparison"):
            gr.Markdown("### Compare predictions across all models")
            gr.Markdown("🔍 **Compare all 5 models at once**: BERT & LogReg for bail outcome, BERT & LogReg & XGBoost for fairness analysis!")
            
            with gr.Row():
                with gr.Column():
                    case_input_compare = gr.Textbox(
                        label="Case Details",
                        placeholder="Enter case details to compare all models... Try the sample cases for comprehensive comparison!",
                        lines=8
                    )
                    
                    with gr.Row():
                        crime_type_compare = gr.Dropdown(
                            choices=CRIME_TYPES,
                            value="Unknown",
                            label="Crime Type"
                        )
                        accused_gender_compare = gr.Dropdown(
                            choices=GENDERS,
                            value="Male",
                            label="Accused Gender"
                        )
                    
                    compare_all_button = gr.Button("Compare All Models", variant="primary", size="lg")
                
                with gr.Column():
                    comparison_results = gr.JSON(label="All Model Results")
    
    # Academic Information Section
    gr.Markdown("""
    ---
    ## 📚 Research Information
    
    ### Research by Sneha Deshmukh and Prathmesh Kamble
    
    This system analyzes Indian bail judgments using multiple machine learning approaches to predict:
    1. **Bail Outcome**: Whether bail will be granted or rejected
    2. **Fairness Analysis**: Whether the judgment shows potential bias
    
    **Models Used**:
    - **LegalBERT**: Fine-tuned BERT model for legal text
    - **Logistic Regression**: Traditional ML with TF-IDF features  
    - **XGBoost**: Gradient boosting for fairness analysis
    
    **Dataset**: [IndianBailJudgments-1200](https://huggingface.co/datasets/SnehaDeshmukh/IndianBailJudgments-1200)  
    **Paper**: [ArXiv:2507.02506](https://arxiv.org/pdf/2507.02506)
    
    **Citation**:
    ```bibtex
    @misc{indianbail2025,
      title  = {IndianBailJudgments-1200: Annotated Dataset of 1200 Indian Bail Judgments},
      author = {Sneha Deshmukh and Prathmesh Kamble},
      year   = {2025},
      url    = {https://huggingface.co/datasets/SnehaDeshmukh/IndianBailJudgments-1200},
      note   = {Dataset publicly released under CC BY 4.0 License}
    }
    ```
    """)

    # Event handlers
    def compare_all_models(case_text, crime_type, accused_gender):
        """Compare predictions from all models"""
        results = {}
        
        # Bail outcome predictions
        results["Bail Outcome - LegalBERT"] = predict_bail_bert(case_text)
        results["Bail Outcome - LogReg"] = predict_bail_logreg(case_text, crime_type, "Regular")
        
        # Fairness predictions
        results["Fairness - LegalBERT"] = predict_fairness_bert(case_text)
        results["Fairness - LogReg"] = predict_fairness_logreg(case_text, accused_gender)
        results["Fairness - XGBoost"] = predict_fairness_xgboost(case_text, accused_gender)
        
        return results
    
    def load_and_populate_sample():
        sample_text1, sample_text2, sample_text3, gender, crime, bail, region = load_sample_case()
        return sample_text1, sample_text2, sample_text3, gender, crime, bail, region
    
    # Bail outcome event handlers
    predict_bert_bail.click(
        predict_bail_bert,
        inputs=[case_input_bail],
        outputs=[result_bert_bail]
    )
    
    predict_logreg_bail.click(
        predict_bail_logreg,
        inputs=[case_input_bail, crime_type_bail, bail_type_bail],
        outputs=[result_logreg_bail]
    )
    
    # Fairness event handlers
    predict_bert_fairness.click(
        predict_fairness_bert,
        inputs=[case_input_fairness],
        outputs=[result_bert_fairness]
    )
    
    predict_logreg_fairness.click(
        predict_fairness_logreg,
        inputs=[case_input_fairness, accused_gender],
        outputs=[result_logreg_fairness]
    )
    
    predict_xgb_fairness.click(
        predict_fairness_xgboost,
        inputs=[case_input_fairness, accused_gender],
        outputs=[result_xgb_fairness]
    )
    
    # Comparison event handler
    compare_all_button.click(
        compare_all_models,
        inputs=[case_input_compare, crime_type_compare, accused_gender_compare],
        outputs=[comparison_results]
    )
    
    # Sample loading event handler
    load_sample_bail.click(
        load_and_populate_sample,
        outputs=[case_input_bail, case_input_fairness, case_input_compare, accused_gender, crime_type_bail, bail_type_bail, region_fairness]
    )

if __name__ == "__main__":
    print("🌐 Starting Gradio app...")
    demo.launch()