Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import numpy as np | |
| from transformers import AutoModel, AutoTokenizer | |
| from datasets import load_dataset | |
| # 1. Rebuild the Brain's Architecture | |
| class LegalExpertModel(nn.Module): | |
| def __init__(self, model_name, num_labels): | |
| super().__init__() | |
| self.bert = AutoModel.from_pretrained(model_name) | |
| self.dropout = nn.Dropout(0.3) | |
| self.classifier = nn.Linear(768, num_labels) | |
| def forward(self, input_ids, attention_mask, **kwargs): | |
| outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, **kwargs) | |
| pooled_output = outputs.last_hidden_state[:, 0, :] | |
| pooled_output = self.dropout(pooled_output) | |
| logits = self.classifier(pooled_output) | |
| return {"logits": logits} | |
| class IPCClassifier: | |
| def __init__(self, model_dir="./expert_ipc_brain"): | |
| print("Loading Expert AI Engine...") | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # --- THE AUTO-TRANSLATOR FIX --- | |
| print("Fetching real IPC names...") | |
| dataset = load_dataset("Exploration-Lab/IL-TUR", "lsi", split='train[:1]', trust_remote_code=True) | |
| real_ipc_names = dataset.features['labels'].feature.names | |
| # Load the internal numbers and translate them to real laws (e.g., "13" -> "IPC 453") | |
| internal_classes = np.load(f"{model_dir}/ipc_classes.npy", allow_pickle=True) | |
| self.classes = [real_ipc_names[int(c)].replace("Section ", "IPC ") if c.isdigit() else c for c in internal_classes] | |
| num_labels = len(self.classes) | |
| # ------------------------------- | |
| # Load the Optimized Math Thresholds | |
| self.thresholds = np.load(f"{model_dir}/optimal_thresholds.npy") | |
| # Load Tokenizer | |
| self.tokenizer = AutoTokenizer.from_pretrained(model_dir) | |
| # Load the Custom Brain Weights | |
| self.model = LegalExpertModel("nlpaueb/legal-bert-base-uncased", num_labels) | |
| self.model.load_state_dict(torch.load(f"{model_dir}/custom_model_weights.bin", map_location=self.device)) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| print("System Ready.") | |
| def predict(self, text_scenario): | |
| """Takes a raw legal string and returns the predicted IPC sections.""" | |
| encodings = self.tokenizer(text_scenario, truncation=True, padding='max_length', max_length=512, return_tensors="pt") | |
| input_ids = encodings['input_ids'].to(self.device) | |
| attention_mask = encodings['attention_mask'].to(self.device) | |
| with torch.no_grad(): | |
| outputs = self.model(input_ids=input_ids, attention_mask=attention_mask) | |
| logits = outputs["logits"].squeeze() | |
| probabilities = torch.sigmoid(logits).cpu().numpy() | |
| predicted_laws = [] | |
| for i, prob in enumerate(probabilities): | |
| if prob > self.thresholds[i]: | |
| predicted_laws.append({ | |
| "ipc_section": self.classes[i], | |
| "confidence": round(float(prob * 100), 2) | |
| }) | |
| if not predicted_laws: | |
| return [{"ipc_section": "No crime detected or insufficient details", "confidence": 0.0}] | |
| return sorted(predicted_laws, key=lambda x: x['confidence'], reverse=True) | |
| # ========================================== | |
| # HOW TO USE IT (Student 1's Web Backend) | |
| # ========================================== | |
| if __name__ == "__main__": | |
| ai_engine = IPCClassifier(model_dir="./expert_ipc_brain") | |
| fake_case = """ | |
| On the night of October 12th, the accused forcefully broke the lock of the | |
| complainant's house and entered without permission. The accused then took | |
| gold jewelry worth Rs. 50,000 from the locker and fled the scene. When | |
| confronted by the guard, the accused brandished a knife and threatened him. | |
| """ | |
| results = ai_engine.predict(fake_case) | |
| print("\n--- AI PREDICTIONS ---") | |
| for res in results: | |
| print(f"Law: {res['ipc_section']} | Confidence: {res['confidence']}%") |