Upload main.py

main.py

@@ -0,0 +1,66 @@
+from transformers import pipeline
+import gradio as gr
+
+# Load pre-trained pipelines
+try:
+    summarizer = pipeline("summarization", model="sshleifer/distilbart-cnn-12-6")
+    ner = pipeline("ner", model="Davlan/bert-base-multilingual-cased-ner-hrl", aggregation_strategy="simple")
+except Exception as e:
+    summarizer = None
+    ner = None
+    print("Error loading models:", e)
+
+# Nigerian law reference (basic keyword-to-punishment mapping)
+crime_punishment_map = {
+    "theft": {"law": "Criminal Code Act, Section 390", "punishment": "Up to 3 years imprisonment"},
+    "armed robbery": {"law": "Robbery and Firearms Act, Section 1", "punishment": "Death penalty or life imprisonment"},
+    "internet fraud": {"law": "Cybercrime Act, 2015", "punishment": "Minimum of 7 years imprisonment"},
+    "rape": {"law": "Criminal Law of Lagos State, Section 260", "punishment": "Life imprisonment"},
+    "murder": {"law": "Criminal Code Act, Section 319", "punishment": "Death penalty"},
+    "assault": {"law": "Criminal Code Act, Section 351", "punishment": "1 year imprisonment"}
+}
+
+def classify_crime(text):
+    text = text.lower()
+    for crime in crime_punishment_map:
+        if crime in text:
+            return crime, crime_punishment_map[crime]
+    return "unknown", {
+        "law": "N/A",
+        "punishment": "No specific punishment found. Manual review required."
+    }
+
+# Main analysis function with full error handling
+def analyze_text(text):
+    try:
+        if not text.strip():
+            return "No text provided.", [], {"Crime Type": "N/A", "Applicable Law": "N/A", "Recommended Punishment": "N/A"}
+
+        summary = summarizer(text, max_length=80, min_length=30, do_sample=False)[0].get("summary_text", "Summary failed.")
+        entities = ner(text)
+        crime_type, law_info = classify_crime(text)
+
+        return summary, entities, {
+            "Crime Type": crime_type.title() if crime_type != "unknown" else "Unknown",
+            "Applicable Law": law_info["law"],
+            "Recommended Punishment": law_info["punishment"]
+        }
+    except Exception as e:
+        return f"⚠️ An error occurred: {str(e)}", [], {
+            "Crime Type": "Error",
+            "Applicable Law": "Error",
+            "Recommended Punishment": "Error"
+        }
+
+# Launch app
+gr.Interface(
+    fn=analyze_text,
+    inputs=gr.Textbox(lines=12, label="Paste Criminal Case Text"),
+    outputs=[
+        gr.Textbox(label="Summary"),
+        gr.JSON(label="Extracted Entities"),
+        gr.JSON(label="Legal Analysis / Recommended Punishment")
+    ],
+    title="JusticeAI - Legal Case Analyzer",
+    description="Paste any criminal case report. This AI will summarize it, extract important entities, and recommend the legal punishment based on Nigerian law."
+).launch()