Spaces:
Sleeping
Sleeping
| import spaces | |
| import os | |
| import torch | |
| import torch.nn as nn | |
| import gradio as gr | |
| from transformers import BigBirdTokenizer, BigBirdModel | |
| from huggingface_hub import hf_hub_download | |
| # --- 1. PyTorch Architecture Definition --- | |
| SCALAR_COLS = [ | |
| "bail_type_mapped", "days_in_custody_available", "days_in_custody", | |
| "age_available_mapped", "min_age", "max_age", "median_age", | |
| "shap_sum_pos", "shap_sum_neg", "shap_max_pos", "shap_min_neg", | |
| "shap_pos_count", "shap_neg_count", "shap_l1_total", | |
| "shap_top3_abs_sum", "xgb_pred" | |
| ] | |
| class BigBirdScalarClassifier(nn.Module): | |
| def __init__(self, bigbird_name="google/bigbird-roberta-large", scalar_dim=16): | |
| super().__init__() | |
| self.bigbird = BigBirdModel.from_pretrained(bigbird_name) | |
| hidden = self.bigbird.config.hidden_size | |
| self.att_vector = nn.Parameter(torch.randn(hidden)) | |
| self.drop_text1 = nn.Dropout(0.1) | |
| self.proj = nn.Linear(hidden, 512) | |
| self.ln = nn.LayerNorm(512) | |
| self.drop_text2 = nn.Dropout(0.2) | |
| self.classifier = nn.Linear(512 + scalar_dim, 1) | |
| def attention_pool(self, H, attention_mask): | |
| scores = torch.matmul(H, self.att_vector) | |
| scores = scores.masked_fill(attention_mask == 0, -1e4) | |
| alpha = torch.softmax(scores, dim=1) | |
| return torch.sum(H * alpha.unsqueeze(-1), dim=1) | |
| def forward(self, input_ids, attention_mask, scalars): | |
| out = self.bigbird(input_ids=input_ids, attention_mask=attention_mask, return_dict=True) | |
| H = out.last_hidden_state | |
| pooled = self.attention_pool(H, attention_mask) | |
| x_text = self.drop_text1(pooled) | |
| x_text = self.proj(x_text) | |
| x_text = self.ln(x_text) | |
| x_text = self.drop_text2(x_text) | |
| x = torch.cat([x_text, scalars], dim=1) | |
| return self.classifier(x).squeeze(1) | |
| # --- 2. Global Model Initialization --- | |
| # ZeroGPU handles the CUDA allocation automatically behind the scenes | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-large") | |
| # Download directly from your HF model repository | |
| model_path = hf_hub_download(repo_id="puspeshsrivastava/ibpsv0.2", filename="best_model_full.pt") | |
| classifier = BigBirdScalarClassifier(scalar_dim=len(SCALAR_COLS)).to(device) | |
| checkpoint = torch.load(model_path, map_location=device) | |
| state_dict = checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint | |
| classifier.load_state_dict(state_dict) | |
| classifier.eval() | |
| # --- 3. Inference Logic (with ZeroGPU Decorator) --- | |
| def predict_bail_outcome( | |
| case_details, bail_type_mapped, days_in_custody_available, days_in_custody, | |
| age_available_mapped, min_age, max_age, median_age, | |
| shap_sum_pos, shap_sum_neg, shap_max_pos, shap_min_neg, | |
| shap_pos_count, shap_neg_count, shap_l1_total, shap_top3_abs_sum, xgb_pred | |
| ): | |
| try: | |
| encoded = tokenizer( | |
| case_details, padding=True, truncation=True, max_length=1200, return_tensors="pt" | |
| ) | |
| input_ids = encoded["input_ids"].to(device) | |
| attention_mask = encoded["attention_mask"].to(device) | |
| scalars_list = [ | |
| float(bail_type_mapped), float(days_in_custody_available), float(days_in_custody), | |
| float(age_available_mapped), float(min_age), float(max_age), float(median_age), | |
| float(shap_sum_pos), float(shap_sum_neg), float(shap_max_pos), float(shap_min_neg), | |
| float(shap_pos_count), float(shap_neg_count), float(shap_l1_total), float(shap_top3_abs_sum), float(xgb_pred) | |
| ] | |
| scalars_tensor = torch.tensor([scalars_list], dtype=torch.float32).to(device) | |
| with torch.no_grad(): | |
| logit = classifier(input_ids, attention_mask, scalars_tensor) | |
| prob = float(torch.sigmoid(logit).item()) | |
| pred_label = 1 if prob >= 0.5 else 0 | |
| return { | |
| "prediction": pred_label, | |
| "probability": round(prob, 4), | |
| "outcome_label": "Bail Granted" if pred_label == 1 else "Bail Denied" | |
| } | |
| except Exception as e: | |
| return {"error": str(e)} | |
| # --- 4. Gradio UI & Auto-API Mapping --- | |
| inputs = [ | |
| gr.Textbox(label="Case Details (Narrative)", lines=4), | |
| gr.Number(label="Bail Type (1=Regular, 2=Anticipatory)", value=1), | |
| gr.Number(label="Days in Custody Available (1=Yes, 0=No)", value=1), | |
| gr.Number(label="Days in Custody", value=0.0), | |
| gr.Number(label="Age Available (1=Yes, 0=No)", value=1), | |
| gr.Number(label="Min Age", value=35.0), | |
| gr.Number(label="Max Age", value=35.0), | |
| gr.Number(label="Median Age", value=35.0), | |
| gr.Number(label="SHAP Sum Pos", value=0.0), | |
| gr.Number(label="SHAP Sum Neg", value=0.0), | |
| gr.Number(label="SHAP Max Pos", value=0.0), | |
| gr.Number(label="SHAP Min Neg", value=0.0), | |
| gr.Number(label="SHAP Pos Count", value=0.0), | |
| gr.Number(label="SHAP Neg Count", value=0.0), | |
| gr.Number(label="SHAP L1 Total", value=0.0), | |
| gr.Number(label="SHAP Top 3 Abs Sum", value=0.0), | |
| gr.Number(label="XGB Pred Baseline", value=0.5), | |
| ] | |
| demo = gr.Interface( | |
| fn=predict_bail_outcome, | |
| inputs=inputs, | |
| outputs=gr.JSON(label="Prediction Response"), | |
| title="IBPS BigBird Predictor (ZeroGPU)", | |
| description="Interactive UI & API for the Indian Bail Prediction System." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |