Johan713 commited on
Commit
d0d6945
·
verified ·
1 Parent(s): b8ac434

Upload app2.py

Browse files
Files changed (1) hide show
  1. app2.py +424 -0
app2.py ADDED
@@ -0,0 +1,424 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from dataclasses import dataclass
3
+ import os
4
+ from uuid import uuid4
5
+ import requests
6
+ import wikipedia
7
+ import googlesearch
8
+ from sentence_transformers import SentenceTransformer
9
+ import PyPDF2
10
+ import docx
11
+ import faiss
12
+ import numpy as np
13
+ import json
14
+ import re
15
+ from sklearn.feature_extraction.text import TfidfVectorizer
16
+ from concurrent.futures import ThreadPoolExecutor
17
+ import nltk
18
+ import spacy
19
+ from dotenv import load_dotenv
20
+
21
+ load_dotenv()
22
+
23
+ nltk.download('wordnet')
24
+ nltk.download('punkt')
25
+
26
+ try:
27
+ spacy.cli.download("en_core_web_sm")
28
+ except Exception as e:
29
+ print(f"Error downloading spacy model: {e}")
30
+
31
+ DEPLOYED = os.getenv("DEPLOYED", "false").lower() == "true"
32
+ MODEL_NAME = "tiiuae/falcon-180B-chat"
33
+ HEADERS = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
34
+ ENDPOINT_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
35
+ DEFAULT_INSTRUCTIONS = """LexAI is an advanced legal AI assistant powered by Falcon 180B, with capabilities including contract analysis, legal research, predictive litigation analysis, intelligent legal drafting, answering common legal questions, document summarization, legal entity recognition, sentiment analysis, and more. LexAI can perform document retrieval, Wikipedia searches, and internet searches to provide comprehensive assistance."""
36
+
37
+ sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
38
+ ner_model = spacy.load("en_core_web_sm")
39
+
40
+ index = faiss.IndexFlatL2(384)
41
+
42
+ tfidf_vectorizer = TfidfVectorizer(stop_words='english')
43
+
44
+ @dataclass
45
+ class Rating:
46
+ prompt: str
47
+ response: str
48
+ ratings: list[str]
49
+
50
+ class Document:
51
+ def __init__(self, content: str, metadata: dict):
52
+ self.content = content
53
+ self.metadata = metadata
54
+ self.embedding = None
55
+
56
+ def compute_embedding(self):
57
+ self.embedding = sentence_model.encode([self.content])[0]
58
+
59
+ class DocumentStore:
60
+ def __init__(self):
61
+ self.documents = []
62
+ self.index = faiss.IndexFlatL2(384)
63
+
64
+ def add_document(self, document: Document):
65
+ document.compute_embedding()
66
+ self.documents.append(document)
67
+ self.index.add(np.array([document.embedding]))
68
+
69
+ def search(self, query: str, k: int = 5):
70
+ query_vector = sentence_model.encode([query])[0]
71
+ distances, indices = self.index.search(np.array([query_vector]), k)
72
+ return [self.documents[i] for i in indices[0]]
73
+
74
+ document_store = DocumentStore()
75
+
76
+ def extract_text_from_file(file):
77
+ if file.name.endswith('.pdf'):
78
+ reader = PyPDF2.PdfReader(file)
79
+ text = ""
80
+ for page in reader.pages:
81
+ text += page.extract_text()
82
+ elif file.name.endswith('.docx'):
83
+ doc = docx.Document(file)
84
+ text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
85
+ else:
86
+ text = file.read().decode('utf-8')
87
+ return text
88
+
89
+ def query_falcon(prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> str:
90
+ payload = {
91
+ "inputs": prompt,
92
+ "parameters": {
93
+ "max_new_tokens": max_tokens,
94
+ "do_sample": True,
95
+ "temperature": temperature,
96
+ "top_p": 0.9,
97
+ "stop": ["User:"]
98
+ }
99
+ }
100
+
101
+ response = requests.post(ENDPOINT_URL, headers=HEADERS, json=payload)
102
+ if response.status_code == 200:
103
+ return response.json()[0]['generated_text']
104
+ else:
105
+ print(f"Error: {response.status_code} - {response.text}")
106
+ return "Error occurred while querying Falcon 180B."
107
+
108
+ def summarize_text(text: str, max_length: int = 150) -> str:
109
+ prompt = f"Summarize the following text in about {max_length} words:\n\n{text}\n\nSummary:"
110
+ return query_falcon(prompt, max_tokens=max_length)
111
+
112
+ def extract_legal_entities(text: str):
113
+ doc = ner_model(text)
114
+ return [ent.text for ent in doc.ents if ent.label_ in ["PERSON", "ORG", "GPE", "LAW"]]
115
+
116
+ def analyze_sentiment(text: str) -> str:
117
+ prompt = f"Analyze the sentiment of the following text. Respond with either 'Positive', 'Negative', or 'Neutral':\n\n{text}\n\nSentiment:"
118
+ return query_falcon(prompt, max_tokens=10)
119
+
120
+ def extract_keywords(text: str, top_n: int = 5):
121
+ prompt = f"Extract the top {top_n} keywords from the following text:\n\n{text}\n\nKeywords:"
122
+ response = query_falcon(prompt, max_tokens=50)
123
+ return [keyword.strip() for keyword in response.split(',')][:top_n]
124
+
125
+ def get_legal_definitions(term: str) -> str:
126
+ prompt = f"Provide a legal definition for the term '{term}':"
127
+ return query_falcon(prompt, max_tokens=100).strip()
128
+
129
+ def perform_case_law_search(query: str):
130
+ prompt = f"Perform a case law search for the following query and provide a summary of relevant cases:\n\n{query}\n\nRelevant cases:"
131
+ return query_falcon(prompt, max_tokens=200)
132
+
133
+ def generate_legal_document(document_type: str, details: dict) -> str:
134
+ prompt = f"Generate a {document_type} with the following details:\n"
135
+ for key, value in details.items():
136
+ prompt += f"{key}: {value}\n"
137
+ prompt += f"\nGenerated {document_type}:"
138
+ return query_falcon(prompt, max_tokens=500).strip()
139
+
140
+ def perform_wikipedia_search(query):
141
+ try:
142
+ search_results = wikipedia.search(query)
143
+ if search_results:
144
+ page = wikipedia.page(search_results[0])
145
+ summary = wikipedia.summary(search_results[0], sentences=3)
146
+ return f"Wikipedia: {summary}\n\nFull article: {page.url}"
147
+ else:
148
+ return "No Wikipedia results found."
149
+ except:
150
+ return "Error occurred while searching Wikipedia."
151
+
152
+ def perform_internet_search(query):
153
+ try:
154
+ search_results = list(googlesearch.search(query, num_results=3))
155
+ if search_results:
156
+ return "Internet search results:\n" + "\n".join(search_results)
157
+ else:
158
+ return "No internet search results found."
159
+ except:
160
+ return "Error occurred while performing internet search."
161
+
162
+ def chat_accordion():
163
+ with gr.Accordion("Parameters", open=False):
164
+ temperature = gr.Slider(
165
+ minimum=0.1,
166
+ maximum=1.0,
167
+ value=0.7,
168
+ step=0.1,
169
+ interactive=True,
170
+ label="Temperature",
171
+ )
172
+ top_p = gr.Slider(
173
+ minimum=0.1,
174
+ maximum=0.99,
175
+ value=0.9,
176
+ step=0.01,
177
+ interactive=True,
178
+ label="p (nucleus sampling)",
179
+ )
180
+
181
+ max_tokens = gr.Slider(
182
+ minimum=64,
183
+ maximum=1024,
184
+ value=64,
185
+ step=1,
186
+ interactive=True,
187
+ label="Max Tokens",
188
+ )
189
+
190
+ session_id = gr.Textbox(
191
+ value=uuid4,
192
+ interactive=False,
193
+ visible=False,
194
+ )
195
+
196
+ with gr.Accordion("Instructions", open=False, visible=False):
197
+ instructions = gr.Textbox(
198
+ placeholder="The Instructions",
199
+ value=DEFAULT_INSTRUCTIONS,
200
+ lines=16,
201
+ interactive=True,
202
+ label="Instructions",
203
+ max_lines=16,
204
+ show_label=False,
205
+ )
206
+ with gr.Row():
207
+ with gr.Column():
208
+ user_name = gr.Textbox(
209
+ lines=1,
210
+ label="username",
211
+ value="User",
212
+ interactive=True,
213
+ placeholder="Username: ",
214
+ show_label=False,
215
+ max_lines=1,
216
+ )
217
+ with gr.Column():
218
+ bot_name = gr.Textbox(
219
+ lines=1,
220
+ value="LexAI",
221
+ interactive=True,
222
+ placeholder="Bot Name",
223
+ show_label=False,
224
+ max_lines=1,
225
+ visible=False,
226
+ )
227
+
228
+ return temperature, top_p, instructions, user_name, bot_name, session_id, max_tokens
229
+
230
+ def format_chat_prompt(message: str, chat_history, instructions: str, user_name: str, bot_name: str):
231
+ instructions = instructions or DEFAULT_INSTRUCTIONS
232
+ instructions = instructions.strip()
233
+ prompt = instructions
234
+ for turn in chat_history:
235
+ user_message, bot_message = turn
236
+ prompt = f"{prompt}\n{user_name}: {user_message}\n{bot_name}: {bot_message}"
237
+ prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
238
+ return prompt
239
+
240
+ def run_chat(message: str, history, instructions: str, user_name: str, bot_name: str, temperature: float, top_p: float, session_id: str, max_tokens: int, uploaded_file: gr.File):
241
+ if uploaded_file is not None:
242
+ document_text = extract_text_from_file(uploaded_file)
243
+ summary = summarize_text(document_text)
244
+ legal_entities = extract_legal_entities(document_text)
245
+ sentiment = analyze_sentiment(document_text)
246
+ keywords = extract_keywords(document_text)
247
+
248
+ doc = Document(content=document_text, metadata={
249
+ "filename": uploaded_file.name,
250
+ "summary": summary,
251
+ "legal_entities": legal_entities,
252
+ "sentiment": sentiment,
253
+ "keywords": keywords
254
+ })
255
+ document_store.add_document(doc)
256
+
257
+ message += f"\n[System: A document '{uploaded_file.name}' has been uploaded and processed.]"
258
+
259
+ relevant_docs = document_store.search(message)
260
+ retrieved_context = "\n".join([doc.content for doc in relevant_docs])
261
+
262
+ with ThreadPoolExecutor(max_workers=2) as executor:
263
+ wiki_future = executor.submit(perform_wikipedia_search, message)
264
+ internet_future = executor.submit(perform_internet_search, message)
265
+
266
+ wiki_result = wiki_future.result()
267
+ internet_result = internet_future.result()
268
+
269
+ case_law_results = perform_case_law_search(message)
270
+
271
+ full_context = f"""Retrieved Documents:\n{retrieved_context}
272
+
273
+ Wikipedia Search:\n{wiki_result}
274
+
275
+ Internet Search:\n{internet_result}
276
+
277
+ Relevant Case Law:\n{case_law_results}
278
+ """
279
+
280
+ prompt = format_chat_prompt(message, history, instructions, user_name, bot_name)
281
+ prompt += f"\nAdditional Context:\n{full_context}\n\nBased on the above information, please provide a comprehensive response:"
282
+
283
+ response = query_falcon(prompt, max_tokens=max_tokens, temperature=temperature)
284
+
285
+ response = post_process_output(response, message)
286
+
287
+ return response
288
+
289
+ def post_process_output(output: str, original_query: str) -> str:
290
+ if "define" in original_query.lower() or "meaning of" in original_query.lower():
291
+ terms = re.findall(r'\b(?!(?:the|a|an)\b)\w+', original_query)
292
+ for term in terms:
293
+ definition = get_legal_definitions(term)
294
+ output += f"\n\nLegal definition of '{term}': {definition}"
295
+
296
+ if "draft" in original_query.lower() or "create document" in original_query.lower():
297
+ doc_type_match = re.search(r'draft a (\w+)', original_query.lower())
298
+ if doc_type_match:
299
+ doc_type = doc_type_match.group(1)
300
+ details = extract_document_details(original_query)
301
+ generated_document = generate_legal_document(doc_type, details)
302
+ output += f"\n\nGenerated {doc_type.capitalize()}:\n\n{generated_document}"
303
+
304
+ if "analyze" in original_query.lower() and "document" in original_query.lower():
305
+ sentiment = analyze_sentiment(output)
306
+ output += f"\n\nOverall sentiment of the analysis: {sentiment}"
307
+
308
+ return output
309
+
310
+ def extract_document_details(query: str) -> dict:
311
+ details = {}
312
+ if "parties" in query.lower():
313
+ details["parties"] = re.search(r'parties: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
314
+ if "date" in query.lower():
315
+ details["date"] = re.search(r'date: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
316
+ if "terms" in query.lower():
317
+ details["terms"] = re.search(r'terms: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
318
+ return details
319
+
320
+ def chat_tab():
321
+ with gr.Column():
322
+ with gr.Row():
323
+ (
324
+ temperature,
325
+ top_p,
326
+ instructions,
327
+ user_name,
328
+ bot_name,
329
+ session_id,
330
+ max_tokens
331
+ ) = chat_accordion()
332
+
333
+ with gr.Column():
334
+ with gr.Blocks():
335
+ prompt_examples = [
336
+ ["Analyze this contract for potential risks and summarize key points."],
337
+ ["Find relevant case law for an intellectual property dispute in the software industry."],
338
+ ["What are the key clauses in a non-disclosure agreement? Draft a template."],
339
+ ["Predict the outcome of this employment discrimination case based on recent precedents."],
340
+ ["Draft a cease and desist letter for trademark infringement. Parties: TechCorp and InnovateNow, Date: 2023-07-22"],
341
+ ]
342
+ file_upload = gr.File(label="Upload Legal Document for Analysis")
343
+ gr.ChatInterface(
344
+ fn=run_chat,
345
+ chatbot=gr.Chatbot(
346
+ height=620,
347
+ render=False,
348
+ show_label=False,
349
+ avatar_images=("images/user_icon.png", "images/lexai_icon.png"),
350
+ ),
351
+ textbox=gr.Textbox(
352
+ placeholder="Ask LexAI about legal matters, analyze documents, or request drafting assistance...",
353
+ render=False,
354
+ scale=7,
355
+ ),
356
+ examples=prompt_examples,
357
+ additional_inputs=[
358
+ instructions,
359
+ user_name,
360
+ bot_name,
361
+ temperature,
362
+ top_p,
363
+ session_id,
364
+ max_tokens,
365
+ file_upload
366
+ ],
367
+ submit_btn="Send",
368
+ stop_btn="Stop",
369
+ retry_btn="🔄 Retry",
370
+ undo_btn="↩️ Delete",
371
+ clear_btn="🗑️ Clear",
372
+ )
373
+
374
+ def introduction():
375
+ with gr.Column(scale=2):
376
+ gr.Image("images/lexai_logo.png", elem_id="banner-image", show_label=False)
377
+ with gr.Column(scale=5):
378
+ gr.Markdown(
379
+ """# LexAI: Advanced Legal AI Assistant
380
+ **LexAI is a cutting-edge AI-driven legal assistant with a wide range of capabilities to support legal professionals and provide information to the public.**
381
+
382
+ ✨ This demo is powered by the state-of-the-art Falcon 180B language model and specialized legal knowledge.
383
+
384
+ 🧪 LexAI offers the following key features:
385
+ 1. Contract Analysis and Risk Assessment
386
+ 2. Legal Research Assistant with Case Law Search
387
+ 3. Predictive Litigation Analysis
388
+ 4. Intelligent Legal Document Drafting
389
+ 5. Legal Chatbot for Public Access
390
+ 6. Document Retrieval, Analysis, and Summarization
391
+ 7. Legal Entity Recognition
392
+ 8. Sentiment Analysis for Legal Texts
393
+ 9. Keyword Extraction from Legal Documents
394
+ 10. Legal Definition Lookup
395
+ 11. Wikipedia and Internet Search Integration
396
+
397
+ ⚠️ **Disclaimer**: LexAI is an AI assistant and does not substitute for professional legal advice. Always consult with a qualified legal professional for specific legal matters.
398
+ """
399
+ )
400
+
401
+ def main():
402
+ with gr.Blocks(
403
+ css="""#chat_container {height: 820px; width: 1000px; margin-left: auto; margin-right: auto;}
404
+ #chatbot {height: 600px; overflow: auto;}
405
+ #create_container {height: 750px; margin-left: 0px; margin-right: 0px;}
406
+ #tokenizer_renderer span {white-space: pre-wrap}
407
+ """
408
+ ) as demo:
409
+ with gr.Row():
410
+ introduction()
411
+ with gr.Row():
412
+ chat_tab()
413
+
414
+ return demo
415
+
416
+ def start_demo():
417
+ demo = main()
418
+ if DEPLOYED:
419
+ demo.queue().launch(show_api=False)
420
+ else:
421
+ demo.queue().launch(share=True)
422
+
423
+ if __name__ == "__main__":
424
+ start_demo()