Spaces:
Sleeping
Sleeping
Upload 12 files
Browse files- .dockerignore +10 -0
- .env +2 -0
- .gitignore +100 -0
- Dockerfile +27 -20
- LICENSE +21 -0
- README.md +19 -19
- api.py +361 -0
- app.py +605 -0
- config.py +34 -0
- requirements.txt +0 -0
- train_model.py +141 -0
- vertex_config.json +13 -0
.dockerignore
ADDED
|
@@ -0,0 +1,10 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
.env
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.pyc
|
| 4 |
+
data/
|
| 5 |
+
notebooks/
|
| 6 |
+
test_*.py
|
| 7 |
+
diag_output*.txt
|
| 8 |
+
.git/
|
| 9 |
+
.gitignore
|
| 10 |
+
venv/
|
.env
ADDED
|
@@ -0,0 +1,2 @@
|
|
|
|
|
|
|
|
|
|
| 1 |
+
|
| 2 |
+
GEMINI_API_KEY=AIzaSyDqLnQkZpNuciaGWLl3yzwVHxxhDXMzcpU
|
.gitignore
ADDED
|
@@ -0,0 +1,100 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Python
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.py[codz]
|
| 4 |
+
*$py.class
|
| 5 |
+
*.so
|
| 6 |
+
.Python
|
| 7 |
+
|
| 8 |
+
# Distribution / packaging
|
| 9 |
+
build/
|
| 10 |
+
develop-eggs/
|
| 11 |
+
dist/
|
| 12 |
+
downloads/
|
| 13 |
+
eggs/
|
| 14 |
+
.eggs/
|
| 15 |
+
lib/
|
| 16 |
+
lib64/
|
| 17 |
+
parts/
|
| 18 |
+
sdist/
|
| 19 |
+
var/
|
| 20 |
+
wheels/
|
| 21 |
+
share/python-wheels/
|
| 22 |
+
*.egg-info/
|
| 23 |
+
.installed.cfg
|
| 24 |
+
*.egg
|
| 25 |
+
MANIFEST
|
| 26 |
+
|
| 27 |
+
# PyInstaller
|
| 28 |
+
*.manifest
|
| 29 |
+
*.spec
|
| 30 |
+
|
| 31 |
+
# Installer logs
|
| 32 |
+
pip-log.txt
|
| 33 |
+
pip-delete-this-directory.txt
|
| 34 |
+
|
| 35 |
+
# Unit test / coverage reports
|
| 36 |
+
htmlcov/
|
| 37 |
+
.tox/
|
| 38 |
+
.nox/
|
| 39 |
+
.coverage
|
| 40 |
+
.coverage.*
|
| 41 |
+
.cache
|
| 42 |
+
nosetests.xml
|
| 43 |
+
coverage.xml
|
| 44 |
+
*.cover
|
| 45 |
+
*.py.cover
|
| 46 |
+
.hypothesis/
|
| 47 |
+
.pytest_cache/
|
| 48 |
+
cover/
|
| 49 |
+
|
| 50 |
+
# Django / Flask / Scrapy
|
| 51 |
+
*.log
|
| 52 |
+
local_settings.py
|
| 53 |
+
db.sqlite3
|
| 54 |
+
db.sqlite3-journal
|
| 55 |
+
instance/
|
| 56 |
+
.webassets-cache
|
| 57 |
+
.scrapy
|
| 58 |
+
|
| 59 |
+
# Notebooks
|
| 60 |
+
.ipynb_checkpoints
|
| 61 |
+
|
| 62 |
+
# Environments & Secrets
|
| 63 |
+
.env
|
| 64 |
+
.envrc
|
| 65 |
+
.venv
|
| 66 |
+
env/
|
| 67 |
+
venv/
|
| 68 |
+
ENV/
|
| 69 |
+
env.bak/
|
| 70 |
+
venv.bak/
|
| 71 |
+
vertex_config.json
|
| 72 |
+
secrets.toml
|
| 73 |
+
*.json
|
| 74 |
+
|
| 75 |
+
# Large Model Files (Upload to HF directly)
|
| 76 |
+
legal_bert_finetuned_risk/
|
| 77 |
+
*.safetensors
|
| 78 |
+
*.bin
|
| 79 |
+
*.h5
|
| 80 |
+
|
| 81 |
+
# OS Files
|
| 82 |
+
.DS_Store
|
| 83 |
+
Thumbs.db
|
| 84 |
+
|
| 85 |
+
# IDEs
|
| 86 |
+
.vscode/
|
| 87 |
+
.idea/
|
| 88 |
+
.spyderproject
|
| 89 |
+
.spyproject
|
| 90 |
+
|
| 91 |
+
# Diagnostic/Log Files
|
| 92 |
+
diag_output*.txt
|
| 93 |
+
|
| 94 |
+
# Cursor / AI
|
| 95 |
+
.cursorignore
|
| 96 |
+
.cursorindexingignore
|
| 97 |
+
__marimo__
|
| 98 |
+
.abstra/
|
| 99 |
+
.mypy_cache/
|
| 100 |
+
.ruff_cache/
|
Dockerfile
CHANGED
|
@@ -1,20 +1,27 @@
|
|
| 1 |
-
|
| 2 |
-
|
| 3 |
-
|
| 4 |
-
|
| 5 |
-
RUN apt-get update && apt-get install -y \
|
| 6 |
-
build-essential \
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
&& rm -rf /var/lib/apt/lists/*
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Use a Python slim image to keep it lightweight
|
| 2 |
+
FROM python:3.10-slim
|
| 3 |
+
|
| 4 |
+
# Install system dependencies for OCR and PDF processing
|
| 5 |
+
RUN apt-get update && apt-get install -y \
|
| 6 |
+
build-essential \
|
| 7 |
+
libgl1-mesa-glx \
|
| 8 |
+
libglib2.0-0 \
|
| 9 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 10 |
+
|
| 11 |
+
# Set working directory
|
| 12 |
+
WORKDIR /app
|
| 13 |
+
|
| 14 |
+
# Copy requirements and install
|
| 15 |
+
# Note: We use --no-cache-dir to keep the image small
|
| 16 |
+
COPY requirements.txt .
|
| 17 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 18 |
+
|
| 19 |
+
# Copy the rest of the application
|
| 20 |
+
COPY . .
|
| 21 |
+
|
| 22 |
+
# Expose Streamlit port
|
| 23 |
+
EXPOSE 8501
|
| 24 |
+
|
| 25 |
+
# Command to run the application
|
| 26 |
+
# We use --server.address=0.0.0.0 for Cloud Run/Container compatibility
|
| 27 |
+
CMD ["streamlit", "run", "app.py", "--server.port=8501", "--server.address=0.0.0.0"]
|
LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2026 khushshah103
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
README.md
CHANGED
|
@@ -1,19 +1,19 @@
|
|
| 1 |
-
---
|
| 2 |
-
title:
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
-
sdk: docker
|
| 7 |
-
app_port: 8501
|
| 8 |
-
|
| 9 |
-
-
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
--
|
| 13 |
-
|
| 14 |
-
#
|
| 15 |
-
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: LegalAI Analyzer
|
| 3 |
+
emoji: ⚖️
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: indigo
|
| 6 |
+
sdk: docker
|
| 7 |
+
app_port: 8501
|
| 8 |
+
pinned: false
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
# LegalAI - AI-powered Contract Intelligence Platform
|
| 12 |
+
AI-powered contract intelligence platform for automated contract review, structured summaries, risk detection and legal entity extraction. Features also include multi-document vendor comparison, conversational document chat and compliance auditing across major industry frameworks.
|
| 13 |
+
|
| 14 |
+
## 🚀 Features
|
| 15 |
+
Professional AI-powered contract review, structured summaries, risk detection, and legal entity extraction.
|
| 16 |
+
|
| 17 |
+
## 🛠️ Local Setup
|
| 18 |
+
1. `pip install -r requirements.txt`
|
| 19 |
+
2. `streamlit run app.py`
|
api.py
ADDED
|
@@ -0,0 +1,361 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
import torch
|
| 3 |
+
import re
|
| 4 |
+
import json
|
| 5 |
+
import time
|
| 6 |
+
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
|
| 7 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 8 |
+
from transformers import pipeline
|
| 9 |
+
from google import genai
|
| 10 |
+
from dotenv import load_dotenv
|
| 11 |
+
|
| 12 |
+
from src.improved_extractor import ImprovedExtractor
|
| 13 |
+
from src.rag_service import RAGService
|
| 14 |
+
from io import BytesIO
|
| 15 |
+
import config
|
| 16 |
+
|
| 17 |
+
# Load environment variables
|
| 18 |
+
load_dotenv()
|
| 19 |
+
|
| 20 |
+
app = FastAPI(title="LegalAI API", version="1.0")
|
| 21 |
+
|
| 22 |
+
# CORS for frontend
|
| 23 |
+
app.add_middleware(
|
| 24 |
+
CORSMiddleware,
|
| 25 |
+
allow_origins=["*"],
|
| 26 |
+
allow_credentials=True,
|
| 27 |
+
allow_methods=["*"],
|
| 28 |
+
allow_headers=["*"],
|
| 29 |
+
)
|
| 30 |
+
|
| 31 |
+
# Global instances
|
| 32 |
+
classifier = None
|
| 33 |
+
client = None
|
| 34 |
+
extractor = None
|
| 35 |
+
rag = None
|
| 36 |
+
|
| 37 |
+
def load_services():
|
| 38 |
+
global classifier, client, extractor, rag
|
| 39 |
+
|
| 40 |
+
# 1. Risk Classifier
|
| 41 |
+
torch.manual_seed(42)
|
| 42 |
+
local_path = config.FINE_TUNED_MODEL_PATH
|
| 43 |
+
model_to_load = "nlpaueb/legal-bert-base-uncased"
|
| 44 |
+
if os.path.isdir(local_path) and os.path.exists(os.path.join(local_path, "config.json")):
|
| 45 |
+
model_to_load = local_path
|
| 46 |
+
|
| 47 |
+
try:
|
| 48 |
+
classifier = pipeline(
|
| 49 |
+
"text-classification",
|
| 50 |
+
model=model_to_load,
|
| 51 |
+
device=0 if torch.cuda.is_available() else -1,
|
| 52 |
+
model_kwargs={"low_cpu_mem_usage": True}
|
| 53 |
+
)
|
| 54 |
+
print(f"✅ API: Classifier loaded using {model_to_load}")
|
| 55 |
+
except Exception as e:
|
| 56 |
+
print(f"⚠️ API: Classifier failed: {e}")
|
| 57 |
+
|
| 58 |
+
# 2. Gemini LLM (Centralized)
|
| 59 |
+
if config.GEMINI_API_KEY:
|
| 60 |
+
try:
|
| 61 |
+
client = genai.Client(api_key=config.GEMINI_API_KEY)
|
| 62 |
+
print("✅ API: Gemini Client initialized")
|
| 63 |
+
except Exception as e:
|
| 64 |
+
print(f"⚠️ API: Gemini Client failed: {e}")
|
| 65 |
+
else:
|
| 66 |
+
# Fallback to Vertex if config exists but key doesn't
|
| 67 |
+
vertex_json = os.path.join(config.BASE_DIR, "vertex_config.json")
|
| 68 |
+
if os.path.exists(vertex_json):
|
| 69 |
+
try:
|
| 70 |
+
with open(vertex_json, "r") as f:
|
| 71 |
+
v_config = json.load(f)
|
| 72 |
+
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(vertex_json)
|
| 73 |
+
client = genai.Client(vertexai=True, project=v_config.get("project_id"), location="us-central1")
|
| 74 |
+
print("✅ API: Vertex AI Client initialized")
|
| 75 |
+
except Exception as e:
|
| 76 |
+
print(f"⚠️ API: Vertex AI failed: {e}")
|
| 77 |
+
|
| 78 |
+
# 3. Extractor
|
| 79 |
+
extractor = ImprovedExtractor()
|
| 80 |
+
|
| 81 |
+
# 4. RAG Service
|
| 82 |
+
rag = RAGService()
|
| 83 |
+
try:
|
| 84 |
+
rag_data_path = os.path.join(config.BASE_DIR, "data", "text")
|
| 85 |
+
if os.path.exists(rag_data_path):
|
| 86 |
+
rag.load_documents(rag_data_path)
|
| 87 |
+
print(f"✅ API: RAG loaded from {rag_data_path}")
|
| 88 |
+
except Exception as e:
|
| 89 |
+
print(f"⚠️ API: RAG loading failed: {e}")
|
| 90 |
+
|
| 91 |
+
@app.on_event("startup")
|
| 92 |
+
async def startup_event():
|
| 93 |
+
load_services()
|
| 94 |
+
|
| 95 |
+
|
| 96 |
+
# ====================================================================
|
| 97 |
+
# Helper: Call Gemini with standardized fallback
|
| 98 |
+
# ====================================================================
|
| 99 |
+
def call_gemini(prompt, preferred_model=None):
|
| 100 |
+
"""Single helper for all Gemini calls. Retries with backoff and model fallback."""
|
| 101 |
+
if not client:
|
| 102 |
+
print("❌ API: No Gemini client available")
|
| 103 |
+
return None
|
| 104 |
+
|
| 105 |
+
# Use preferred model if provided, else fallback to config list
|
| 106 |
+
models_to_try = [preferred_model] if preferred_model else config.MODEL_FALLBACKS
|
| 107 |
+
|
| 108 |
+
for attempt in range(2): # 2 main retry loops
|
| 109 |
+
for model_id in models_to_try:
|
| 110 |
+
try:
|
| 111 |
+
# Basic generation
|
| 112 |
+
response = client.models.generate_content(model=model_id, contents=prompt)
|
| 113 |
+
if response and response.text:
|
| 114 |
+
return response.text
|
| 115 |
+
else:
|
| 116 |
+
print(f"⚠️ API: {model_id} returned empty response")
|
| 117 |
+
except Exception as e:
|
| 118 |
+
err = str(e).upper()
|
| 119 |
+
# Handle Rate Limits
|
| 120 |
+
if "429" in err or "QUOTA" in err or "LIMIT" in err:
|
| 121 |
+
wait = 10 * (attempt + 1)
|
| 122 |
+
print(f"⏳ API: Rate limited on {model_id}. Waiting {wait}s...")
|
| 123 |
+
time.sleep(wait)
|
| 124 |
+
break # Try next model or next attempt
|
| 125 |
+
# Handle Model Not Found (404)
|
| 126 |
+
elif "404" in err or "NOT FOUND" in err:
|
| 127 |
+
print(f"❌ API: Model {model_id} NOT FOUND. Skipping.")
|
| 128 |
+
continue # Try next model in list
|
| 129 |
+
else:
|
| 130 |
+
print(f"❌ API: {model_id} failed: {e}")
|
| 131 |
+
continue
|
| 132 |
+
return None
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
# ====================================================================
|
| 136 |
+
# Utility
|
| 137 |
+
# ====================================================================
|
| 138 |
+
def clean_text(text):
|
| 139 |
+
if not text:
|
| 140 |
+
return ""
|
| 141 |
+
text = re.sub(r'_{2,}', '', text)
|
| 142 |
+
text = re.sub(r'\{.*?\}', '', text)
|
| 143 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 144 |
+
text = re.sub(r' +', ' ', text)
|
| 145 |
+
text = re.sub(r'Page \d+', '', text)
|
| 146 |
+
return text.strip()
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
# ====================================================================
|
| 150 |
+
# API Models
|
| 151 |
+
# ====================================================================
|
| 152 |
+
from pydantic import BaseModel
|
| 153 |
+
from typing import List, Optional
|
| 154 |
+
|
| 155 |
+
class TextRequest(BaseModel):
|
| 156 |
+
text: str
|
| 157 |
+
|
| 158 |
+
class ChatRequest(BaseModel):
|
| 159 |
+
text: str
|
| 160 |
+
history: List[dict]
|
| 161 |
+
prompt: str
|
| 162 |
+
|
| 163 |
+
class ComplianceRequest(BaseModel):
|
| 164 |
+
text: str
|
| 165 |
+
framework: str
|
| 166 |
+
|
| 167 |
+
class SearchRequest(BaseModel):
|
| 168 |
+
query: str
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
# ====================================================================
|
| 172 |
+
# ENDPOINTS (each defined ONCE, clean and simple)
|
| 173 |
+
# ====================================================================
|
| 174 |
+
|
| 175 |
+
@app.get("/api/health")
|
| 176 |
+
async def health_check():
|
| 177 |
+
return {
|
| 178 |
+
"status": "healthy",
|
| 179 |
+
"services": {
|
| 180 |
+
"classifier": classifier is not None,
|
| 181 |
+
"gemini": client is not None,
|
| 182 |
+
"rag": rag is not None
|
| 183 |
+
}
|
| 184 |
+
}
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
@app.post("/api/analyze/extract")
|
| 188 |
+
async def extract_document(file: UploadFile = File(...)):
|
| 189 |
+
if not file:
|
| 190 |
+
raise HTTPException(status_code=400, detail="No file provided")
|
| 191 |
+
temp_path = f"temp_{file.filename}"
|
| 192 |
+
try:
|
| 193 |
+
content = await file.read()
|
| 194 |
+
with open(temp_path, "wb") as f:
|
| 195 |
+
f.write(content)
|
| 196 |
+
raw_text = extractor.extract_text(temp_path)
|
| 197 |
+
cleaned_text = clean_text(raw_text)
|
| 198 |
+
return {"raw_text": raw_text, "cleaned_text": cleaned_text}
|
| 199 |
+
finally:
|
| 200 |
+
if os.path.exists(temp_path):
|
| 201 |
+
os.remove(temp_path)
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
@app.post("/api/analyze/summary")
|
| 205 |
+
async def analyze_summary(req: TextRequest):
|
| 206 |
+
prompt = f"""
|
| 207 |
+
Act as a Senior Legal Counsel. Provide a professional, detailed, and structured executive summary.
|
| 208 |
+
Use bold headings for: Purpose & Overview, Key Obligations, Payment & Compensation, Term & Termination, and Liability & Risk.
|
| 209 |
+
|
| 210 |
+
Contract Text:
|
| 211 |
+
{req.text[:60000]}
|
| 212 |
+
"""
|
| 213 |
+
result = call_gemini(prompt)
|
| 214 |
+
return {"summary": result or "⚠️ Summary generation failed."}
|
| 215 |
+
|
| 216 |
+
|
| 217 |
+
@app.post("/api/analyze/entities")
|
| 218 |
+
async def analyze_entities(req: TextRequest):
|
| 219 |
+
# Smart sampling: first 20k + last 15k to capture preamble AND signature blocks
|
| 220 |
+
first_part = req.text[:20000]
|
| 221 |
+
last_part = req.text[-15000:] if len(req.text) > 20000 else ""
|
| 222 |
+
sample = first_part + "\n\n--- END OF DOCUMENT ---\n\n" + last_part
|
| 223 |
+
|
| 224 |
+
prompt = f"""
|
| 225 |
+
Act as a Legal Clerk. Extract core legal entities from this contract.
|
| 226 |
+
|
| 227 |
+
CRITICAL RULES:
|
| 228 |
+
- Do NOT extract placeholders like [PROVIDER LEGAL NAME] or [CUSTOMER].
|
| 229 |
+
- Check the signature block at the end for actual names.
|
| 230 |
+
- If only placeholders exist, write "NOT SPECIFIED (Generic Template)".
|
| 231 |
+
|
| 232 |
+
Extract:
|
| 233 |
+
1. Contracting Parties (Full legal names)
|
| 234 |
+
2. Effective Date
|
| 235 |
+
3. Governing Law
|
| 236 |
+
4. Total Contract Value
|
| 237 |
+
|
| 238 |
+
Contract Text:
|
| 239 |
+
{sample}
|
| 240 |
+
"""
|
| 241 |
+
result = call_gemini(prompt)
|
| 242 |
+
return {"entities_text": result or "⚠️ Entity extraction failed."}
|
| 243 |
+
|
| 244 |
+
|
| 245 |
+
@app.post("/api/analyze/scam")
|
| 246 |
+
async def analyze_scam(req: TextRequest):
|
| 247 |
+
prompt = f"""
|
| 248 |
+
Act as a Senior Contract Auditor. Scan for predatory, hidden, or highly imbalanced clauses.
|
| 249 |
+
Focus on: IP transfers, uncapped liability, sneaky auto-renewals, hidden exit fees.
|
| 250 |
+
|
| 251 |
+
If found, respond: FLAGGED: [1-sentence explanation]
|
| 252 |
+
If safe, respond: SAFE
|
| 253 |
+
|
| 254 |
+
Contract Text:
|
| 255 |
+
{req.text[:60000]}
|
| 256 |
+
"""
|
| 257 |
+
result = call_gemini(prompt)
|
| 258 |
+
if result and "FLAGGED:" in result:
|
| 259 |
+
return {"scam_warning": result.split("FLAGGED:")[1].strip()}
|
| 260 |
+
return {"scam_warning": None}
|
| 261 |
+
|
| 262 |
+
|
| 263 |
+
@app.post("/api/analyze/risk")
|
| 264 |
+
async def analyze_risk(req: TextRequest):
|
| 265 |
+
if not classifier:
|
| 266 |
+
return {"label": "N/A", "score": 0.0, "description": ""}
|
| 267 |
+
try:
|
| 268 |
+
cleaned = clean_text(req.text)
|
| 269 |
+
result = classifier(cleaned[:512])[0]
|
| 270 |
+
label_id = result['label']
|
| 271 |
+
mapping = {
|
| 272 |
+
"LABEL_0": ("High Risk", "Critical issues found. Requires legal review."),
|
| 273 |
+
"LABEL_1": ("Low Risk", "Standard safe clauses. Low legal overhead."),
|
| 274 |
+
"LABEL_2": ("Medium Risk", "Minor deviations found. Proceed with caution.")
|
| 275 |
+
}
|
| 276 |
+
name, desc = mapping.get(label_id, (label_id, ""))
|
| 277 |
+
return {"label": name, "score": result['score'], "description": desc}
|
| 278 |
+
except Exception as e:
|
| 279 |
+
return {"label": f"Error: {e}", "score": 0.0, "description": ""}
|
| 280 |
+
|
| 281 |
+
|
| 282 |
+
@app.post("/api/analyze/compliance")
|
| 283 |
+
async def analyze_compliance(req: ComplianceRequest):
|
| 284 |
+
prompt = f"""
|
| 285 |
+
Act as an expert compliance auditor. Check this contract against: '{req.framework}'.
|
| 286 |
+
|
| 287 |
+
Evaluate 4-5 critical requirements. For each, give Pass (✅) or Fail (❌) with a 1-sentence reason.
|
| 288 |
+
Format as a clean Markdown list.
|
| 289 |
+
|
| 290 |
+
Contract Text:
|
| 291 |
+
{req.text[:60000]}
|
| 292 |
+
"""
|
| 293 |
+
result = call_gemini(prompt)
|
| 294 |
+
return {"result": result or "⚠️ Compliance check failed."}
|
| 295 |
+
|
| 296 |
+
|
| 297 |
+
@app.post("/api/analyze/compare")
|
| 298 |
+
async def analyze_compare(req: TextRequest):
|
| 299 |
+
prompt = f"""
|
| 300 |
+
Extract the following strictly as a JSON object with EXACTLY these keys:
|
| 301 |
+
"Vendor_Name", "Total_Pricing", "Term_Duration", "Liability_Cap", "Termination_Notice"
|
| 302 |
+
|
| 303 |
+
If a field is missing, use "Not Specified".
|
| 304 |
+
Do NOT output any markdown blocks. ONLY output the raw JSON object.
|
| 305 |
+
|
| 306 |
+
Contract text:
|
| 307 |
+
{req.text[:15000]}
|
| 308 |
+
"""
|
| 309 |
+
result = call_gemini(prompt)
|
| 310 |
+
if result:
|
| 311 |
+
try:
|
| 312 |
+
resp_text = result.strip()
|
| 313 |
+
if resp_text.startswith("```json"):
|
| 314 |
+
resp_text = resp_text[7:]
|
| 315 |
+
if resp_text.endswith("```"):
|
| 316 |
+
resp_text = resp_text[:-3]
|
| 317 |
+
data = json.loads(resp_text)
|
| 318 |
+
return {"data": data}
|
| 319 |
+
except Exception:
|
| 320 |
+
pass
|
| 321 |
+
return {"data": None}
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
@app.post("/api/chat")
|
| 325 |
+
async def chat_document(req: ChatRequest):
|
| 326 |
+
history_text = ""
|
| 327 |
+
for msg in req.history:
|
| 328 |
+
role_str = "User" if msg["role"] == "user" else "Assistant"
|
| 329 |
+
history_text += f"{role_str}: {msg['content']}\n"
|
| 330 |
+
|
| 331 |
+
prompt = f"""
|
| 332 |
+
You are a helpful legal assistant. Answer based ONLY on the contract text.
|
| 333 |
+
If the answer is not in the text, say "I cannot find the answer to this in the document."
|
| 334 |
+
|
| 335 |
+
Contract Text:
|
| 336 |
+
{req.text[:100000]}
|
| 337 |
+
|
| 338 |
+
Conversation History:
|
| 339 |
+
{history_text}
|
| 340 |
+
|
| 341 |
+
Latest Question: {req.prompt}
|
| 342 |
+
"""
|
| 343 |
+
result = call_gemini(prompt)
|
| 344 |
+
return {"answer": result or "Failed to generate answer."}
|
| 345 |
+
|
| 346 |
+
|
| 347 |
+
@app.post("/api/library/search")
|
| 348 |
+
async def library_search(req: SearchRequest):
|
| 349 |
+
if not rag:
|
| 350 |
+
return {"error": "RAG service not loaded"}
|
| 351 |
+
try:
|
| 352 |
+
relevant_chunks = rag.query(req.query, top_k=3)
|
| 353 |
+
answer = rag.generate_answer(req.query, relevant_chunks, client)
|
| 354 |
+
refs = [{"file": res["metadata"]["file"], "text": res["text"]} for res in relevant_chunks]
|
| 355 |
+
return {"answer": answer, "references": refs}
|
| 356 |
+
except Exception as e:
|
| 357 |
+
return {"error": f"Search failed: {e}"}
|
| 358 |
+
|
| 359 |
+
if __name__ == "__main__":
|
| 360 |
+
import uvicorn
|
| 361 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
app.py
ADDED
|
@@ -0,0 +1,605 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import streamlit as st
|
| 2 |
+
import os
|
| 3 |
+
import torch
|
| 4 |
+
from transformers import pipeline
|
| 5 |
+
from src.improved_extractor import ImprovedExtractor
|
| 6 |
+
from src.rag_service import RAGService
|
| 7 |
+
import re
|
| 8 |
+
import json
|
| 9 |
+
import time
|
| 10 |
+
import pandas as pd
|
| 11 |
+
from google import genai
|
| 12 |
+
from dotenv import load_dotenv
|
| 13 |
+
import config
|
| 14 |
+
|
| 15 |
+
# Load environment variables
|
| 16 |
+
load_dotenv()
|
| 17 |
+
|
| 18 |
+
# --- Page Configuration ---
|
| 19 |
+
st.set_page_config(
|
| 20 |
+
page_title=config.APP_NAME,
|
| 21 |
+
page_icon="⚖️",
|
| 22 |
+
layout="wide",
|
| 23 |
+
)
|
| 24 |
+
|
| 25 |
+
# Initialize Session State Globally with robust checks
|
| 26 |
+
def init_session_state():
|
| 27 |
+
if "messages" not in st.session_state:
|
| 28 |
+
st.session_state["messages"] = []
|
| 29 |
+
if "current_file" not in st.session_state:
|
| 30 |
+
st.session_state["current_file"] = None
|
| 31 |
+
if "file_data" not in st.session_state:
|
| 32 |
+
st.session_state["file_data"] = {}
|
| 33 |
+
if "models_loaded" not in st.session_state:
|
| 34 |
+
st.session_state["models_loaded"] = False
|
| 35 |
+
|
| 36 |
+
init_session_state()
|
| 37 |
+
|
| 38 |
+
# --- Styling ---
|
| 39 |
+
st.markdown("""
|
| 40 |
+
<style>
|
| 41 |
+
.main {
|
| 42 |
+
background-color: #f5f7f9;
|
| 43 |
+
}
|
| 44 |
+
.stButton>button {
|
| 45 |
+
width: 100%;
|
| 46 |
+
border-radius: 5px;
|
| 47 |
+
height: 3em;
|
| 48 |
+
background-color: #1e3a8a;
|
| 49 |
+
color: white;
|
| 50 |
+
}
|
| 51 |
+
</style>
|
| 52 |
+
""", unsafe_allow_html=True)
|
| 53 |
+
|
| 54 |
+
# --- Global Services ---
|
| 55 |
+
@st.cache_resource(show_spinner=False)
|
| 56 |
+
def get_classifier(version=1):
|
| 57 |
+
# Fix seed for consistent random head initialization (if fine-tuned path is missing)
|
| 58 |
+
torch.manual_seed(42)
|
| 59 |
+
|
| 60 |
+
# Check for local fine-tuned model from the user's notebook progress
|
| 61 |
+
local_path = config.FINE_TUNED_MODEL_PATH
|
| 62 |
+
model_to_load = "nlpaueb/legal-bert-base-uncased"
|
| 63 |
+
|
| 64 |
+
if os.path.isdir(local_path) and os.path.exists(os.path.join(local_path, "config.json")):
|
| 65 |
+
model_to_load = local_path
|
| 66 |
+
print(f"✅ Loading local fine-tuned model: {local_path}")
|
| 67 |
+
else:
|
| 68 |
+
print(f"ℹ️ Falling back to base Legal-BERT (random accuracy for risk).")
|
| 69 |
+
|
| 70 |
+
return pipeline(
|
| 71 |
+
"text-classification",
|
| 72 |
+
model=model_to_load,
|
| 73 |
+
device=0 if torch.cuda.is_available() else -1,
|
| 74 |
+
model_kwargs={"low_cpu_mem_usage": True}
|
| 75 |
+
)
|
| 76 |
+
|
| 77 |
+
@st.cache_resource
|
| 78 |
+
def get_gemini_client():
|
| 79 |
+
# Centralized client initialization
|
| 80 |
+
if config.GEMINI_API_KEY:
|
| 81 |
+
try:
|
| 82 |
+
c = genai.Client(api_key=config.GEMINI_API_KEY)
|
| 83 |
+
print("✅ Initialized Gemini Client")
|
| 84 |
+
return c
|
| 85 |
+
except Exception as e:
|
| 86 |
+
print(f"❌ Gemini Client init failed: {e}")
|
| 87 |
+
|
| 88 |
+
# Fallback to Vertex
|
| 89 |
+
vertex_json = os.path.join(config.BASE_DIR, "vertex_config.json")
|
| 90 |
+
if os.path.exists(vertex_json):
|
| 91 |
+
try:
|
| 92 |
+
with open(vertex_json, "r") as f:
|
| 93 |
+
v_config = json.load(f)
|
| 94 |
+
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = os.path.abspath(vertex_json)
|
| 95 |
+
c = genai.Client(
|
| 96 |
+
vertexai=True,
|
| 97 |
+
project=v_config.get("project_id"),
|
| 98 |
+
location="us-central1"
|
| 99 |
+
)
|
| 100 |
+
print("✅ Initialized Vertex AI Client")
|
| 101 |
+
return c
|
| 102 |
+
except Exception as e:
|
| 103 |
+
print(f"⚠️ Vertex AI failed: {e}")
|
| 104 |
+
|
| 105 |
+
return None
|
| 106 |
+
|
| 107 |
+
# Removed BERT-based get_ner to save RAM. Using Gemini for Entity Extraction.
|
| 108 |
+
|
| 109 |
+
@st.cache_resource
|
| 110 |
+
def get_rag():
|
| 111 |
+
r = RAGService()
|
| 112 |
+
try:
|
| 113 |
+
rag_data_path = os.path.join(config.BASE_DIR, "data", "text")
|
| 114 |
+
if os.path.exists(rag_data_path):
|
| 115 |
+
r.load_documents(rag_data_path)
|
| 116 |
+
except Exception as e:
|
| 117 |
+
print(f"Error loading RAG: {e}")
|
| 118 |
+
return r
|
| 119 |
+
|
| 120 |
+
# --- UI Loader (Non-Cached) ---
|
| 121 |
+
def load_models(progress_bar=None, status_text=None):
|
| 122 |
+
"""
|
| 123 |
+
Orchestrates the loading of all models with UI progress feedback.
|
| 124 |
+
The individual load functions are cached, so this function run fast after the first time.
|
| 125 |
+
"""
|
| 126 |
+
# 1. Risk Classifier
|
| 127 |
+
if status_text: status_text.text("Loading Risk Classifier (BERT)...")
|
| 128 |
+
if progress_bar: progress_bar.progress(20)
|
| 129 |
+
classifier = get_classifier()
|
| 130 |
+
|
| 131 |
+
# 2. Gemini LLM (for Summarization, RAG & NER)
|
| 132 |
+
if status_text: status_text.text("Connecting to Gemini AI...")
|
| 133 |
+
if progress_bar: progress_bar.progress(50)
|
| 134 |
+
client = get_gemini_client()
|
| 135 |
+
|
| 136 |
+
# 3. RAG Service
|
| 137 |
+
if status_text: status_text.text("Indexing Document Library...")
|
| 138 |
+
if progress_bar: progress_bar.progress(90)
|
| 139 |
+
rag_service = get_rag()
|
| 140 |
+
|
| 141 |
+
return {
|
| 142 |
+
"classifier": classifier,
|
| 143 |
+
"client": client,
|
| 144 |
+
"extractor": ImprovedExtractor(),
|
| 145 |
+
"rag": rag_service
|
| 146 |
+
}
|
| 147 |
+
|
| 148 |
+
# --- Initial Setup & Startup Sequence ---
|
| 149 |
+
def startup_sequence():
|
| 150 |
+
"""
|
| 151 |
+
Shows a splash screen only on the first run of the session.
|
| 152 |
+
"""
|
| 153 |
+
if not st.session_state.get('models_loaded'):
|
| 154 |
+
placeholder = st.empty()
|
| 155 |
+
with placeholder.container():
|
| 156 |
+
st.markdown(f"""
|
| 157 |
+
<div style="text-align: center; padding: 50px;">
|
| 158 |
+
<h1 style="font-size: 3em;">⚖️ {config.APP_NAME}</h1>
|
| 159 |
+
<p style="font-size: 1.2em; color: #666;">Warming up the AI engine... This takes about 30-45 seconds on first boot.</p>
|
| 160 |
+
</div>
|
| 161 |
+
""", unsafe_allow_html=True)
|
| 162 |
+
|
| 163 |
+
progress_bar = st.progress(0)
|
| 164 |
+
status_text = st.empty()
|
| 165 |
+
|
| 166 |
+
# This triggers the individual cached loads
|
| 167 |
+
load_models(progress_bar, status_text)
|
| 168 |
+
|
| 169 |
+
progress_bar.progress(100)
|
| 170 |
+
status_text.text("System ready!")
|
| 171 |
+
st.session_state['models_loaded'] = True
|
| 172 |
+
time.sleep(1)
|
| 173 |
+
placeholder.empty()
|
| 174 |
+
|
| 175 |
+
# Run startup
|
| 176 |
+
startup_sequence()
|
| 177 |
+
models = load_models() # Returns instantly if already cached
|
| 178 |
+
classifier = models["classifier"]
|
| 179 |
+
client = models["client"]
|
| 180 |
+
extractor = models["extractor"]
|
| 181 |
+
rag = models["rag"]
|
| 182 |
+
|
| 183 |
+
def call_gemini(prompt, preferred_model=None):
|
| 184 |
+
"""Single helper for all Gemini calls. Retries with backoff and model fallback."""
|
| 185 |
+
if not client:
|
| 186 |
+
return None
|
| 187 |
+
|
| 188 |
+
models_to_try = [preferred_model] if preferred_model else config.MODEL_FALLBACKS
|
| 189 |
+
|
| 190 |
+
for attempt in range(2):
|
| 191 |
+
for model_id in models_to_try:
|
| 192 |
+
try:
|
| 193 |
+
response = client.models.generate_content(model=model_id, contents=prompt)
|
| 194 |
+
if response and response.text:
|
| 195 |
+
return response.text
|
| 196 |
+
except Exception as e:
|
| 197 |
+
err = str(e).upper()
|
| 198 |
+
if "429" in err or "QUOTA" in err or "LIMIT" in err:
|
| 199 |
+
if attempt < 1:
|
| 200 |
+
wait = 10 * (attempt + 1)
|
| 201 |
+
time.sleep(wait)
|
| 202 |
+
break
|
| 203 |
+
continue
|
| 204 |
+
elif "404" in err or "NOT FOUND" in err:
|
| 205 |
+
continue
|
| 206 |
+
else:
|
| 207 |
+
continue
|
| 208 |
+
return None
|
| 209 |
+
|
| 210 |
+
def get_summary(text):
|
| 211 |
+
prompt = f"""
|
| 212 |
+
Act as a Senior Legal Counsel with 20 years of experience in contract law.
|
| 213 |
+
Review the provided legal contract and generate a high-level, professional executive summary.
|
| 214 |
+
|
| 215 |
+
STRUCTURE YOUR RESPONSE AS FOLLOWS:
|
| 216 |
+
1. **Executive Overview**: High-level purpose of the agreement.
|
| 217 |
+
2. **Key Financial Terms**: Payment schedules, amounts, and late fees.
|
| 218 |
+
3. **Operational Obligations**: What must each party actually DO?
|
| 219 |
+
4. **Termination & Exit**: How do parties leave, and what are the notice periods?
|
| 220 |
+
5. **Critical Liability & Risk**: Indemnities, liability caps, and any "one-sided" clauses.
|
| 221 |
+
6. **Counsel's Recommendation**: A 2-3 sentence professional verdict on the contract's fairness.
|
| 222 |
+
|
| 223 |
+
INSTRUCTIONS:
|
| 224 |
+
- Use a professional, objective, and analytical tone.
|
| 225 |
+
- If you find highly imbalanced or "predatory" clauses, mention them under "Critical Liability & Risk" in a factual, legal manner rather than using alarmist language.
|
| 226 |
+
- Focus on specificities (dates, percentages, dollar amounts).
|
| 227 |
+
- Ensure the summary is readable but dense with information.
|
| 228 |
+
|
| 229 |
+
Contract Text:
|
| 230 |
+
{text[:config.MAX_CHAR_LIMIT]}
|
| 231 |
+
"""
|
| 232 |
+
result = call_gemini(prompt)
|
| 233 |
+
return result or "⚠️ Summary generation failed."
|
| 234 |
+
|
| 235 |
+
def get_entities(text):
|
| 236 |
+
"""Smart sampling: first 20k + last 15k chars to capture preamble AND signature blocks."""
|
| 237 |
+
first_part = text[:20000]
|
| 238 |
+
last_part = text[-15000:] if len(text) > 20000 else ""
|
| 239 |
+
sample = first_part + "\n\n--- END OF DOCUMENT ---\n\n" + last_part
|
| 240 |
+
|
| 241 |
+
prompt = f"""
|
| 242 |
+
Act as a Legal Clerk. Extract the following core entities from the contract.
|
| 243 |
+
|
| 244 |
+
CRITICAL ACCURACY RULES:
|
| 245 |
+
1. DO NOT extract generic placeholders like "[PROVIDER LEGAL NAME]", "[CUSTOMER]", or placeholders in curly brackets.
|
| 246 |
+
2. Examine BOTH the introductory paragraph AND the signature blocks at the end for ACTUAL company names.
|
| 247 |
+
3. If a field only contains a placeholder, write "NOT SPECIFIED (Generic Template Detected)".
|
| 248 |
+
|
| 249 |
+
IDENTIFY:
|
| 250 |
+
1. Contracting Parties (Full legal names of all parties involved)
|
| 251 |
+
2. Effective Date (The start date of the agreement)
|
| 252 |
+
3. Governing Law (Which state/country's laws apply)
|
| 253 |
+
4. Total Contract Value (Specific monetary amount or fee structure)
|
| 254 |
+
|
| 255 |
+
Contract Text:
|
| 256 |
+
{sample}
|
| 257 |
+
"""
|
| 258 |
+
result = call_gemini(prompt)
|
| 259 |
+
return result or "⚠️ Entity extraction failed."
|
| 260 |
+
|
| 261 |
+
def check_unethical_clauses(text):
|
| 262 |
+
"""Scans full document (60k chars) for predatory clauses."""
|
| 263 |
+
prompt = f"""
|
| 264 |
+
Act as a Senior Contract Auditor. Scan for predatory, hidden, or highly imbalanced clauses.
|
| 265 |
+
Focus on: IP transfers, uncapped liability, sneaky auto-renewals, hidden exit fees.
|
| 266 |
+
|
| 267 |
+
If found, respond: FLAGGED: [1-sentence explanation]
|
| 268 |
+
If safe, respond: SAFE
|
| 269 |
+
|
| 270 |
+
Contract Text:
|
| 271 |
+
{text[:60000]}
|
| 272 |
+
"""
|
| 273 |
+
result = call_gemini(prompt)
|
| 274 |
+
if result and "FLAGGED:" in result:
|
| 275 |
+
return result.split("FLAGGED:")[1].strip()
|
| 276 |
+
return None
|
| 277 |
+
|
| 278 |
+
|
| 279 |
+
def get_risk(text):
|
| 280 |
+
if not classifier:
|
| 281 |
+
return "N/A", 0.0, ""
|
| 282 |
+
try:
|
| 283 |
+
cleaned = clean_text(text)
|
| 284 |
+
# Use first 512 tokens as BERT limit
|
| 285 |
+
result = classifier(cleaned[:512])[0]
|
| 286 |
+
label_id = result['label']
|
| 287 |
+
mapping = {
|
| 288 |
+
"LABEL_0": ("High Risk", "Critical issues found. Requires legal review."),
|
| 289 |
+
"LABEL_1": ("Low Risk", "Standard safe clauses. Low legal overhead."),
|
| 290 |
+
"LABEL_2": ("Medium Risk", "Minor deviations found. Proceed with caution.")
|
| 291 |
+
}
|
| 292 |
+
name, desc = mapping.get(label_id, (label_id, ""))
|
| 293 |
+
return name, result['score'], desc
|
| 294 |
+
except Exception as e:
|
| 295 |
+
return f"Error: {e}", 0.0, ""
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
|
| 299 |
+
def run_compliance_check(text, framework):
|
| 300 |
+
"""Runs compliance audit against the full document (60k chars)."""
|
| 301 |
+
prompt = f"""
|
| 302 |
+
Act as an expert compliance auditor. Check this contract against: '{framework}'.
|
| 303 |
+
|
| 304 |
+
Evaluate 4-5 critical requirements. For each, give Pass (✅) or Fail (❌) with 1-sentence reason.
|
| 305 |
+
Format as a clean Markdown list.
|
| 306 |
+
|
| 307 |
+
Contract Text:
|
| 308 |
+
{text[:60000]}
|
| 309 |
+
"""
|
| 310 |
+
result = call_gemini(prompt)
|
| 311 |
+
return result or "⚠️ Compliance check failed."
|
| 312 |
+
|
| 313 |
+
def clean_text(text):
|
| 314 |
+
"""
|
| 315 |
+
Cleans raw document text by removing noise like placeholders and extra whitespace.
|
| 316 |
+
"""
|
| 317 |
+
if not text:
|
| 318 |
+
return ""
|
| 319 |
+
# Remove long sequences of underscores (placeholders)
|
| 320 |
+
text = re.sub(r'_{2,}', '', text)
|
| 321 |
+
# Remove curly bracket placeholders e.g. {services/project name}
|
| 322 |
+
text = re.sub(r'\{.*?\}', '', text)
|
| 323 |
+
# Remove multiple newlines
|
| 324 |
+
text = re.sub(r'\n{3,}', '\n\n', text)
|
| 325 |
+
# Remove multiple spaces
|
| 326 |
+
text = re.sub(r' +', ' ', text)
|
| 327 |
+
# Remove "Page X" noise
|
| 328 |
+
text = re.sub(r'Page \d+', '', text)
|
| 329 |
+
return text.strip()
|
| 330 |
+
def get_comparison_data(text):
|
| 331 |
+
prompt = f"""
|
| 332 |
+
Extract the following strictly as a JSON object with EXACTLY these keys:
|
| 333 |
+
"Vendor_Name", "Total_Pricing", "Term_Duration", "Liability_Cap", "Termination_Notice"
|
| 334 |
+
|
| 335 |
+
If a field is missing, use "Not Specified".
|
| 336 |
+
Do NOT output any markdown blocks. ONLY output the raw JSON object.
|
| 337 |
+
|
| 338 |
+
Contract text:
|
| 339 |
+
{text[:15000]}
|
| 340 |
+
"""
|
| 341 |
+
result = call_gemini(prompt)
|
| 342 |
+
if result:
|
| 343 |
+
try:
|
| 344 |
+
resp_text = result.strip()
|
| 345 |
+
if resp_text.startswith("```json"):
|
| 346 |
+
resp_text = resp_text[7:]
|
| 347 |
+
if resp_text.endswith("```"):
|
| 348 |
+
resp_text = resp_text[:-3]
|
| 349 |
+
return json.loads(resp_text)
|
| 350 |
+
except Exception:
|
| 351 |
+
pass
|
| 352 |
+
return None
|
| 353 |
+
|
| 354 |
+
# Check for model health
|
| 355 |
+
is_fine_tuned = os.path.isdir("legal_bert_finetuned_risk")
|
| 356 |
+
|
| 357 |
+
# --- Sidebar ---
|
| 358 |
+
with st.sidebar:
|
| 359 |
+
st.image("https://cdn-icons-png.flaticon.com/512/2901/2901306.png", width=100)
|
| 360 |
+
st.title("Admin Panel")
|
| 361 |
+
st.info("Upload your legal contracts to begin automated analysis. Upload multiple files for vendor comparison.")
|
| 362 |
+
uploaded_files = st.file_uploader("Upload PDF Contract(s)", type=["pdf"], accept_multiple_files=True)
|
| 363 |
+
|
| 364 |
+
# --- Header ---
|
| 365 |
+
st.title("⚖️ Legal Document & Risk Analyzer")
|
| 366 |
+
st.markdown("Automated intelligence for contract review, risk mitigation, and semantic search.")
|
| 367 |
+
|
| 368 |
+
if not is_fine_tuned:
|
| 369 |
+
st.warning("⚠️ **Warning:** No fine-tuned model found in `./legal_bert_finetuned_risk`. The risk classification is currently using 'Base' weights and will be inaccurate. Please provide your trained model files for accurate assessment.")
|
| 370 |
+
|
| 371 |
+
if uploaded_files:
|
| 372 |
+
if len(uploaded_files) == 1:
|
| 373 |
+
uploaded_file = uploaded_files[0]
|
| 374 |
+
# Check if a new file was uploaded to reset the cache
|
| 375 |
+
if st.session_state["current_file"] != uploaded_file.name:
|
| 376 |
+
st.session_state["current_file"] = uploaded_file.name
|
| 377 |
+
st.session_state["messages"] = [] # Reset chat
|
| 378 |
+
st.session_state["file_data"] = {} # Reset analytics cache
|
| 379 |
+
|
| 380 |
+
# Save uploaded file to temp path
|
| 381 |
+
with open(uploaded_file.name, "wb") as f:
|
| 382 |
+
f.write(uploaded_file.getbuffer())
|
| 383 |
+
|
| 384 |
+
# Text Extraction
|
| 385 |
+
with st.spinner("Processing document..."):
|
| 386 |
+
raw_text = extractor.extract_text(uploaded_file.name)
|
| 387 |
+
cleaned_text = clean_text(raw_text)
|
| 388 |
+
|
| 389 |
+
# Clean up temp file immediately so it doesn't clutter the folder
|
| 390 |
+
try:
|
| 391 |
+
os.remove(uploaded_file.name)
|
| 392 |
+
except OSError:
|
| 393 |
+
pass
|
| 394 |
+
|
| 395 |
+
# Phase 1: Summary & Scam Scan
|
| 396 |
+
col_sum, col_risk = st.columns([2, 1])
|
| 397 |
+
|
| 398 |
+
with col_sum:
|
| 399 |
+
st.subheader("🤖 AI Summary")
|
| 400 |
+
if "summary" not in st.session_state["file_data"]:
|
| 401 |
+
with st.spinner("Analyzing document with Gemini 2.5 Pro..."):
|
| 402 |
+
st.session_state["file_data"]["summary"] = get_summary(raw_text)
|
| 403 |
+
st.markdown(st.session_state["file_data"]["summary"])
|
| 404 |
+
|
| 405 |
+
# Note: Scam/Ethical check is still performed but information is now integrated into the Risk and Summary sections
|
| 406 |
+
if "scam_warning" not in st.session_state["file_data"]:
|
| 407 |
+
time.sleep(2) # Rate limit protection
|
| 408 |
+
with st.spinner("Auditing clauses..."):
|
| 409 |
+
st.session_state["file_data"]["scam_warning"] = check_unethical_clauses(raw_text)
|
| 410 |
+
|
| 411 |
+
with col_risk:
|
| 412 |
+
st.subheader("⚖️ Risk Profile")
|
| 413 |
+
if "risk" not in st.session_state["file_data"]:
|
| 414 |
+
with st.spinner("Analyzing risk..."):
|
| 415 |
+
st.session_state["file_data"]["risk"] = get_risk(raw_text)
|
| 416 |
+
|
| 417 |
+
label, score, description = st.session_state["file_data"]["risk"]
|
| 418 |
+
color = "#16a34a" # Low
|
| 419 |
+
if "Medium" in label: color = "#f59e0b"
|
| 420 |
+
if "High" in label: color = "#ef4444"
|
| 421 |
+
|
| 422 |
+
st.markdown(f"""
|
| 423 |
+
<div style="background-color: {color}; padding: 15px; border-radius: 8px; color: white; text-align: center;">
|
| 424 |
+
<h2 style="color: white; margin: 0; font-size: 1.5em;">{label}</h2>
|
| 425 |
+
</div>
|
| 426 |
+
""", unsafe_allow_html=True)
|
| 427 |
+
if description:
|
| 428 |
+
st.markdown(f"**Actionable Insight:** {description}")
|
| 429 |
+
if not is_fine_tuned:
|
| 430 |
+
st.caption("🚨 Results are uncalibrated (Base Model).")
|
| 431 |
+
|
| 432 |
+
# Phase 2: Professional Entity Extraction
|
| 433 |
+
st.divider()
|
| 434 |
+
st.subheader("🔍 Key Legal Entities")
|
| 435 |
+
if "entities" not in st.session_state["file_data"]:
|
| 436 |
+
time.sleep(2) # Rate limit protection
|
| 437 |
+
with st.spinner("Extracting parties with AI..."):
|
| 438 |
+
entities_text = get_entities(raw_text) # Use raw_text for full context
|
| 439 |
+
st.session_state["file_data"]["entities"] = entities_text if entities_text else "No entities detected."
|
| 440 |
+
|
| 441 |
+
st.info(st.session_state["file_data"]["entities"])
|
| 442 |
+
|
| 443 |
+
|
| 444 |
+
# Phase 2.5: Automated Compliance Checklists
|
| 445 |
+
st.divider()
|
| 446 |
+
st.subheader("🛡️ Compliance & Audit")
|
| 447 |
+
st.markdown("Run automated checks against strict regulatory and industry frameworks.")
|
| 448 |
+
|
| 449 |
+
frameworks = [
|
| 450 |
+
"Select a framework to audit...",
|
| 451 |
+
"General Data Protection Regulation (GDPR)",
|
| 452 |
+
"Standard SaaS Agreement Best Practices",
|
| 453 |
+
"Independent Contractor / Freelance Standard"
|
| 454 |
+
]
|
| 455 |
+
|
| 456 |
+
selected_framework = st.selectbox("Select Compliance Framework:", frameworks)
|
| 457 |
+
|
| 458 |
+
if selected_framework != "Select a framework to audit...":
|
| 459 |
+
cache_key = f"compliance_{selected_framework}"
|
| 460 |
+
|
| 461 |
+
if cache_key not in st.session_state["file_data"]:
|
| 462 |
+
with st.spinner(f"Running {selected_framework} audit..."):
|
| 463 |
+
if not client:
|
| 464 |
+
st.warning("Gemini AI is not connected.")
|
| 465 |
+
else:
|
| 466 |
+
st.session_state["file_data"][cache_key] = run_compliance_check(cleaned_text, selected_framework)
|
| 467 |
+
|
| 468 |
+
if st.session_state["file_data"].get(cache_key):
|
| 469 |
+
st.info(st.session_state["file_data"][cache_key])
|
| 470 |
+
elif client:
|
| 471 |
+
st.error("Audit failed to generate.")
|
| 472 |
+
|
| 473 |
+
# Phase 3: Interactive Chatbot
|
| 474 |
+
st.divider()
|
| 475 |
+
st.subheader("💬 Chat with Document")
|
| 476 |
+
|
| 477 |
+
if not client:
|
| 478 |
+
st.warning("⚠️ **Gemini API Key Missing:** Interactive chat is disabled.")
|
| 479 |
+
else:
|
| 480 |
+
# Display chat messages from history on app rerun
|
| 481 |
+
for message in st.session_state.messages:
|
| 482 |
+
with st.chat_message(message["role"]):
|
| 483 |
+
st.markdown(message["content"])
|
| 484 |
+
|
| 485 |
+
# Accept user input
|
| 486 |
+
if prompt := st.chat_input("Ask a question about this contract (e.g., 'What are the termination conditions?')..."):
|
| 487 |
+
# Add user message to chat history
|
| 488 |
+
st.session_state.messages.append({"role": "user", "content": prompt})
|
| 489 |
+
# Display user message in chat message container
|
| 490 |
+
with st.chat_message("user"):
|
| 491 |
+
st.markdown(prompt)
|
| 492 |
+
|
| 493 |
+
# Display assistant response in chat message container
|
| 494 |
+
with st.chat_message("assistant"):
|
| 495 |
+
message_placeholder = st.empty()
|
| 496 |
+
with st.spinner("Analyzing document and generating answer..."):
|
| 497 |
+
context = cleaned_text[:100000] # Safe limit for large documents
|
| 498 |
+
|
| 499 |
+
# Build conversation history for the prompt
|
| 500 |
+
history_text = ""
|
| 501 |
+
for msg in st.session_state.messages[:-1]: # exclude the current prompt
|
| 502 |
+
role_str = "User" if msg["role"] == "user" else "Assistant"
|
| 503 |
+
history_text += f"{role_str}: {msg['content']}\n"
|
| 504 |
+
|
| 505 |
+
full_prompt = f"""
|
| 506 |
+
You are a helpful legal assistant. Answer the user's latest question based ONLY on the following contract text and the conversation history so far.
|
| 507 |
+
If the answer is not in the text, say "I cannot find the answer to this in the document."
|
| 508 |
+
|
| 509 |
+
Contract Text:
|
| 510 |
+
{context}
|
| 511 |
+
|
| 512 |
+
Conversation History:
|
| 513 |
+
{history_text}
|
| 514 |
+
|
| 515 |
+
Latest Question: {prompt}
|
| 516 |
+
"""
|
| 517 |
+
|
| 518 |
+
answer = call_gemini(full_prompt)
|
| 519 |
+
if not answer:
|
| 520 |
+
answer = "Failed to generate answer. Please try again."
|
| 521 |
+
message_placeholder.markdown(answer)
|
| 522 |
+
|
| 523 |
+
# Add assistant response to chat history
|
| 524 |
+
st.session_state.messages.append({"role": "assistant", "content": answer})
|
| 525 |
+
else:
|
| 526 |
+
# Multi-Document Vendor Comparison Mode
|
| 527 |
+
st.header("📊 Multi-Document Vendor Comparison Matrix")
|
| 528 |
+
st.markdown("Comparing key terms across multiple uploaded contracts.")
|
| 529 |
+
|
| 530 |
+
comparison_results = []
|
| 531 |
+
|
| 532 |
+
# We use a progress bar to show extraction status
|
| 533 |
+
progress_bar = st.progress(0)
|
| 534 |
+
status_text = st.empty()
|
| 535 |
+
|
| 536 |
+
for i, file in enumerate(uploaded_files):
|
| 537 |
+
status_text.text(f"Extracting data from {file.name} ({i+1}/{len(uploaded_files)})...")
|
| 538 |
+
# Save temp
|
| 539 |
+
with open(file.name, "wb") as f:
|
| 540 |
+
f.write(file.getbuffer())
|
| 541 |
+
|
| 542 |
+
# Extract
|
| 543 |
+
raw_text = extractor.extract_text(file.name)
|
| 544 |
+
cleaned_text = clean_text(raw_text)
|
| 545 |
+
|
| 546 |
+
# Clean up temp file immediately
|
| 547 |
+
try:
|
| 548 |
+
os.remove(file.name)
|
| 549 |
+
except OSError:
|
| 550 |
+
pass
|
| 551 |
+
|
| 552 |
+
# Get structured data
|
| 553 |
+
data = get_comparison_data(cleaned_text)
|
| 554 |
+
if data:
|
| 555 |
+
data["Filename"] = file.name
|
| 556 |
+
comparison_results.append(data)
|
| 557 |
+
else:
|
| 558 |
+
comparison_results.append({
|
| 559 |
+
"Filename": file.name,
|
| 560 |
+
"Vendor_Name": "Extraction Failed",
|
| 561 |
+
"Total_Pricing": "N/A",
|
| 562 |
+
"Term_Duration": "N/A",
|
| 563 |
+
"Liability_Cap": "N/A",
|
| 564 |
+
"Termination_Notice": "N/A"
|
| 565 |
+
})
|
| 566 |
+
|
| 567 |
+
progress_bar.progress((i + 1) / len(uploaded_files))
|
| 568 |
+
|
| 569 |
+
status_text.text("Extraction Complete!")
|
| 570 |
+
|
| 571 |
+
if comparison_results:
|
| 572 |
+
st.divider()
|
| 573 |
+
df = pd.DataFrame(comparison_results)
|
| 574 |
+
# Reorder columns to put Filename first
|
| 575 |
+
cols = ["Filename", "Vendor_Name", "Total_Pricing", "Term_Duration", "Liability_Cap", "Termination_Notice"]
|
| 576 |
+
df = df[[c for c in cols if c in df.columns]]
|
| 577 |
+
|
| 578 |
+
st.dataframe(df, use_container_width=True, hide_index=True)
|
| 579 |
+
|
| 580 |
+
st.success("Comparison Matrix generated successfully! You can download this table via the download button inside the table view.")
|
| 581 |
+
|
| 582 |
+
else:
|
| 583 |
+
st.warning("Please upload a PDF document in the sidebar to start analysis.")
|
| 584 |
+
st.info("💡 **Tip:** You can use the Library Search below to query your repository of legal documents.")
|
| 585 |
+
|
| 586 |
+
st.divider()
|
| 587 |
+
st.subheader("📚 Global Library Search")
|
| 588 |
+
lib_query = st.text_input("Identify patterns across your entire library:")
|
| 589 |
+
if lib_query:
|
| 590 |
+
with st.spinner("Searching library..."):
|
| 591 |
+
relevant_chunks = rag.query(lib_query, top_k=3)
|
| 592 |
+
answer = rag.generate_answer(lib_query, relevant_chunks, client)
|
| 593 |
+
|
| 594 |
+
|
| 595 |
+
st.markdown("### 🤖 Synthesized Knowledge")
|
| 596 |
+
st.success(answer)
|
| 597 |
+
|
| 598 |
+
st.markdown("#### Document References")
|
| 599 |
+
for i, res in enumerate(relevant_chunks):
|
| 600 |
+
st.markdown(f"**{i+1}. From {res['metadata']['file']}:**")
|
| 601 |
+
st.caption(res['text'])
|
| 602 |
+
|
| 603 |
+
# --- Footer ---
|
| 604 |
+
st.divider()
|
| 605 |
+
st.caption("LegalAI Analyzer v1.2 | Powered by Legal-BERT & T5")
|
config.py
ADDED
|
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import os
|
| 2 |
+
from dotenv import load_dotenv
|
| 3 |
+
|
| 4 |
+
# Load environment variables
|
| 5 |
+
load_dotenv()
|
| 6 |
+
|
| 7 |
+
# --- Gemini Configuration ---
|
| 8 |
+
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
|
| 9 |
+
|
| 10 |
+
# Verified models based on diagnostic output
|
| 11 |
+
# Using 2.0-flash as primary for speed, pro-latest for complex tasks
|
| 12 |
+
PRIMARY_MODEL = "gemini-2.0-flash"
|
| 13 |
+
COMPLEX_MODEL = "gemini-pro-latest"
|
| 14 |
+
|
| 15 |
+
# Fallback chain for reliability
|
| 16 |
+
MODEL_FALLBACKS = [
|
| 17 |
+
"gemini-2.0-flash",
|
| 18 |
+
"gemini-flash-latest",
|
| 19 |
+
"gemini-pro-latest",
|
| 20 |
+
"gemini-2.5-flash" # Experimental but available
|
| 21 |
+
]
|
| 22 |
+
|
| 23 |
+
# --- App Settings ---
|
| 24 |
+
APP_NAME = "LegalAI Portable"
|
| 25 |
+
MAX_CHAR_LIMIT = 60000 # Streamlit/API context window limit
|
| 26 |
+
ENTITY_EXTRACTION_CHARS = 35000 # Smart sampling limit
|
| 27 |
+
|
| 28 |
+
# --- Paths ---
|
| 29 |
+
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
| 30 |
+
DATA_DIR = os.path.join(BASE_DIR, "data")
|
| 31 |
+
FINE_TUNED_MODEL_PATH = os.path.join(BASE_DIR, "legal_bert_finetuned_risk")
|
| 32 |
+
|
| 33 |
+
def is_model_loaded():
|
| 34 |
+
return GEMINI_API_KEY is not None and GEMINI_API_KEY != ""
|
requirements.txt
CHANGED
|
Binary files a/requirements.txt and b/requirements.txt differ
|
|
|
train_model.py
ADDED
|
@@ -0,0 +1,141 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import pandas as pd
|
| 2 |
+
import numpy as np
|
| 3 |
+
import torch
|
| 4 |
+
import os
|
| 5 |
+
from datasets import Dataset, DatasetDict, ClassLabel, Features
|
| 6 |
+
from transformers import (
|
| 7 |
+
AutoTokenizer,
|
| 8 |
+
AutoModelForSequenceClassification,
|
| 9 |
+
TrainingArguments,
|
| 10 |
+
Trainer,
|
| 11 |
+
pipeline
|
| 12 |
+
)
|
| 13 |
+
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score
|
| 14 |
+
|
| 15 |
+
# --- 1. Load Dataset ---
|
| 16 |
+
csv_path = os.path.join("data", "legal_contract_clauses.csv")
|
| 17 |
+
if not os.path.exists(csv_path):
|
| 18 |
+
print(f"Error: '{csv_path}' not found locally.")
|
| 19 |
+
import sys
|
| 20 |
+
sys.exit(1)
|
| 21 |
+
|
| 22 |
+
print("Step 1: Loading dataset...")
|
| 23 |
+
df = pd.read_csv(csv_path)
|
| 24 |
+
TEXT_COLUMN = "clause_text"
|
| 25 |
+
LABEL_COLUMN = "risk_level"
|
| 26 |
+
|
| 27 |
+
# --- 2. Create Label Mappings ---
|
| 28 |
+
print("Step 2: Creating label mappings...")
|
| 29 |
+
labels = df[LABEL_COLUMN].unique()
|
| 30 |
+
labels.sort()
|
| 31 |
+
label2id = {label: i for i, label in enumerate(labels)}
|
| 32 |
+
id2label = {i: label for i, label in enumerate(labels)}
|
| 33 |
+
num_labels = len(labels)
|
| 34 |
+
df['label'] = df[LABEL_COLUMN].map(label2id)
|
| 35 |
+
|
| 36 |
+
# --- 3. Convert to Hugging Face Dataset ---
|
| 37 |
+
print("Step 3: Converting to Hugging Face Dataset...")
|
| 38 |
+
dataset = Dataset.from_pandas(df)
|
| 39 |
+
|
| 40 |
+
print("Step 3a: Casting 'label' column to ClassLabel for stratification...")
|
| 41 |
+
label_names_in_order = [id2label[i] for i in range(num_labels)]
|
| 42 |
+
class_label_feature = ClassLabel(names=label_names_in_order)
|
| 43 |
+
dataset = dataset.cast_column("label", class_label_feature)
|
| 44 |
+
|
| 45 |
+
print("Step 3b: Splitting dataset...")
|
| 46 |
+
train_test_split = dataset.train_test_split(test_size=0.2, seed=42, stratify_by_column="label")
|
| 47 |
+
dataset_dict = DatasetDict({
|
| 48 |
+
'train': train_test_split['train'],
|
| 49 |
+
'test': train_test_split['test']
|
| 50 |
+
})
|
| 51 |
+
print(f"Training data: {len(dataset_dict['train'])} examples")
|
| 52 |
+
print(f"Test data: {len(dataset_dict['test'])} examples")
|
| 53 |
+
|
| 54 |
+
# --- 4. Loading Model ---
|
| 55 |
+
print("\nStep 4: Loading model and tokenizer (nlpaueb/legal-bert-base-uncased)...")
|
| 56 |
+
model_name = "nlpaueb/legal-bert-base-uncased"
|
| 57 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
| 58 |
+
model = AutoModelForSequenceClassification.from_pretrained(
|
| 59 |
+
model_name,
|
| 60 |
+
num_labels=num_labels,
|
| 61 |
+
id2label=id2label,
|
| 62 |
+
label2id=label2id
|
| 63 |
+
)
|
| 64 |
+
|
| 65 |
+
# --- 5. Tokenizing ---
|
| 66 |
+
print("\nStep 5: Tokenizing dataset...")
|
| 67 |
+
def tokenize_function(examples):
|
| 68 |
+
return tokenizer(examples[TEXT_COLUMN], padding="max_length", truncation=True)
|
| 69 |
+
|
| 70 |
+
tokenized_datasets = dataset_dict.map(tokenize_function, batched=True)
|
| 71 |
+
columns_to_remove = [TEXT_COLUMN, LABEL_COLUMN]
|
| 72 |
+
if "__index_level_0__" in tokenized_datasets["train"].column_names:
|
| 73 |
+
columns_to_remove.append("__index_level_0__")
|
| 74 |
+
|
| 75 |
+
tokenized_datasets = tokenized_datasets.remove_columns(columns_to_remove)
|
| 76 |
+
tokenized_datasets.set_format("torch")
|
| 77 |
+
print("Tokenization complete.")
|
| 78 |
+
|
| 79 |
+
# --- 6. Metrics ---
|
| 80 |
+
def compute_metrics(eval_pred):
|
| 81 |
+
logits, labels = eval_pred
|
| 82 |
+
predictions = np.argmax(logits, axis=-1)
|
| 83 |
+
acc = accuracy_score(labels, predictions)
|
| 84 |
+
f1 = f1_score(labels, predictions, average='weighted')
|
| 85 |
+
precision = precision_score(labels, predictions, average='weighted')
|
| 86 |
+
recall = recall_score(labels, predictions, average='weighted')
|
| 87 |
+
return {
|
| 88 |
+
"accuracy": acc,
|
| 89 |
+
"f1": f1,
|
| 90 |
+
"precision": precision,
|
| 91 |
+
"recall": recall
|
| 92 |
+
}
|
| 93 |
+
|
| 94 |
+
# --- 7. Set Training Arguments & Train ---
|
| 95 |
+
print("\nStep 6: Setting training arguments...")
|
| 96 |
+
model_output_dir = "./legal_bert_finetuned_risk"
|
| 97 |
+
training_args = TrainingArguments(
|
| 98 |
+
output_dir=model_output_dir,
|
| 99 |
+
learning_rate=2e-5,
|
| 100 |
+
per_device_train_batch_size=8,
|
| 101 |
+
per_device_eval_batch_size=8,
|
| 102 |
+
num_train_epochs=3,
|
| 103 |
+
weight_decay=0.01,
|
| 104 |
+
load_best_model_at_end=True,
|
| 105 |
+
push_to_hub=False,
|
| 106 |
+
report_to="none",
|
| 107 |
+
eval_strategy="epoch",
|
| 108 |
+
save_strategy="epoch",
|
| 109 |
+
)
|
| 110 |
+
|
| 111 |
+
trainer = Trainer(
|
| 112 |
+
model=model,
|
| 113 |
+
args=training_args,
|
| 114 |
+
train_dataset=tokenized_datasets["train"],
|
| 115 |
+
eval_dataset=tokenized_datasets["test"],
|
| 116 |
+
tokenizer=tokenizer,
|
| 117 |
+
compute_metrics=compute_metrics,
|
| 118 |
+
)
|
| 119 |
+
|
| 120 |
+
print("\nStep 7: Starting model training...")
|
| 121 |
+
trainer.train()
|
| 122 |
+
print("Training finished.")
|
| 123 |
+
|
| 124 |
+
# --- 8. Save ---
|
| 125 |
+
trainer.save_model(model_output_dir)
|
| 126 |
+
tokenizer.save_pretrained(model_output_dir)
|
| 127 |
+
print(f"Model saved to {model_output_dir}")
|
| 128 |
+
|
| 129 |
+
# --- 9. Inference Test ---
|
| 130 |
+
print("\nStep 8: Quick inference test...")
|
| 131 |
+
risk_classifier = pipeline(
|
| 132 |
+
"text-classification",
|
| 133 |
+
model=model_output_dir,
|
| 134 |
+
device=0 if torch.cuda.is_available() else -1
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
+
test_clause = "Indemnification. The Contractor agrees to indemnify..."
|
| 138 |
+
result = risk_classifier(test_clause, return_all_scores=True)
|
| 139 |
+
print(f"Test Clause: '{test_clause}'")
|
| 140 |
+
print(f"Prediction: {result}")
|
| 141 |
+
|
vertex_config.json
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"type": "service_account",
|
| 3 |
+
"project_id": "gen-lang-client-0878631680",
|
| 4 |
+
"private_key_id": "d6f22485bc2ec5a3ab7b79452d52122a8ba07bd9",
|
| 5 |
+
"private_key": "-----BEGIN PRIVATE KEY-----\nMIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDMkeS+GQtVFKfG\n2EsznAHhPUlw6iwpLtFRQGjNyKQeMk7g/6CYXZPHNddUZIdWB7fpoVW4tvoUW9uH\nB+Dxs+fHBdqxzd20rsVPjnQlEhqWDwCbzICzn/j1gmaMfmFoR2d8GpZZk/WZLHUx\nqTVvJcelo3o+D2Y1JKcAPP6RSAsu2VF0MSzs0xYCpRokEMeTAaRh8LdIzwn4x2yO\nkG2FpgAvVwNu4C7yqPKqX/G6+47RpDOLE7XxrA7KJQ3Iyaq1izlB/AJv6TpJYCnJ\nF+y0+3ZtG+YWARJKryRYoIM6I3uYbT0hl4aSRrlyL/24BiOMHXMh8Ly/RrGPrlTD\nP2q3jGVhAgMBAAECggEAO+oc9ypJneUbUItH1zlBebEkAdysC3HJX4VATMkPfEKa\nxJ5J8GYz4nwb8X0yABnpnRUiEKQSsYTH0pAVq2TYJPvLdhkH7qPjaS4dSUA4piuF\nr5vhM/rOBUXoeAyJxetn3TrCP0EtBOw29NEMs916UEKFioijFDyEZvb2TiGuGM2J\noS9DCgjpy8khJlF8qRaJXYhRY4IZLOZ7/N6W0v0urBTlw6B678pK3vNNWTGetDUw\nNw9eunHJx0laMo7dPi0ZRKXVU46HjPh+UjO9JkVnSDtQ17nIi+uy1KQk3DM9zSNB\nLZ1S8OMNQ/w7nof5x49KtL4U/hkVBAJEZE+XIbrKEwKBgQD9rC4WCQMBvQQN8v/b\nT24F/Tn+OW3W0nD6MZk2D4T69uJIZoKg9wtcGPSfB/eWGbyrTUl45aYoPPzofAuA\nJtdYm8BBHom3Wl2c7hm7yD+9JriTovEPR2YTAPRb8WT9Uum5660cyeZrXdKUBw/k\nuwHfcu0h0q+QbE5Bx0HgFzd/swKBgQDOcmHfowPPQZh5Apxo6XmQsfro/wItQux6\nVYQPorp0nZc3QuK39TzD0w5EB0ln+PiePNk9Ka2QzksyKibwXDeK6+fFfLiD+VNe\nlhQ9ScQSc7g7aLKpUVgyBUhU3a94C3Fo6VcjD0UaEiGFXnaysjsz0NObuNWjoK/6\ndoMsg9ucmwKBgDQiV9Jlhb2suBGZ2GWdPHJ0qZ+K8/0LgPaajO9kXyMo7DaPAtfM\nDgSeiF/KxmKN5Y2bM7dqEyz/48Zv//kVgnqOgAOiIBGTu9uNv7ItJJReSd7lxP8r\n4FmVf6MJyISbrrdsLUbWY8m7HZUYonaSzH684ulpoHVhHPA94FcKzngrAoGAR5kh\n2pbFtnaQc791EWmPsKAZXf86+QmRzCemYmnUcqdJD3BSAuy5g0zKUMuaKMYPK4n8\n6ZpvyJ3TNHWsBYZl+Pcx7asArzykLlSsdkkECEY76da4x8IfbVGpsU0lEuQdE/6B\nKvabx5vsJj4JupXXtHPDP+mNpm8POTUq7yBX5OsCgYEAzDF3Wa5lLD9aDKfWjgYt\nGcXV+As1S+Lo0Vz4fspDZkwGGJk4Wfe2GIrPIbN+w2ybVlpFY9AskmlyEfveXT77\nKKjqhpdwOG2mN+ZQAbtga6I7x187uzihkMjgnuaEKBMWqSkxbIrDEy2YpJmyQvha\nY3bHgWSRtEwbysSaHCLbDhA=\n-----END PRIVATE KEY-----\n",
|
| 6 |
+
"client_email": "vertex-express@gen-lang-client-0878631680.iam.gserviceaccount.com",
|
| 7 |
+
"client_id": "117615514449694666186",
|
| 8 |
+
"auth_uri": "https://accounts.google.com/o/oauth2/auth",
|
| 9 |
+
"token_uri": "https://oauth2.googleapis.com/token",
|
| 10 |
+
"auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs",
|
| 11 |
+
"client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/vertex-express%40gen-lang-client-0878631680.iam.gserviceaccount.com",
|
| 12 |
+
"universe_domain": "googleapis.com"
|
| 13 |
+
}
|