khushshah103 commited on
Commit
52775bf
·
verified ·
1 Parent(s): c302d39

Upload 12 files

Browse files
Files changed (12) hide show
  1. .dockerignore +10 -0
  2. .env +2 -0
  3. .gitignore +100 -0
  4. Dockerfile +27 -20
  5. LICENSE +21 -0
  6. README.md +19 -19
  7. api.py +361 -0
  8. app.py +605 -0
  9. config.py +34 -0
  10. requirements.txt +0 -0
  11. train_model.py +141 -0
  12. 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
- FROM python:3.13.5-slim
2
-
3
- WORKDIR /app
4
-
5
- RUN apt-get update && apt-get install -y \
6
- build-essential \
7
- curl \
8
- git \
9
- && rm -rf /var/lib/apt/lists/*
10
-
11
- COPY requirements.txt ./
12
- COPY src/ ./src/
13
-
14
- RUN pip3 install -r requirements.txt
15
-
16
- EXPOSE 8501
17
-
18
- HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
19
-
20
- ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
 
 
 
 
 
 
 
 
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: LegalAi
3
- emoji: 🚀
4
- colorFrom: red
5
- colorTo: red
6
- sdk: docker
7
- app_port: 8501
8
- tags:
9
- - streamlit
10
- pinned: false
11
- short_description: Streamlit template space
12
- ---
13
-
14
- # Welcome to Streamlit!
15
-
16
- Edit `/src/streamlit_app.py` to customize this app to your heart's desire. :heart:
17
-
18
- If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
19
- forums](https://discuss.streamlit.io).
 
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
+ }