|
import gradio as gr
|
|
from dataclasses import dataclass
|
|
import os
|
|
from uuid import uuid4
|
|
import requests
|
|
import wikipedia
|
|
import googlesearch
|
|
from sentence_transformers import SentenceTransformer
|
|
import PyPDF2
|
|
import docx
|
|
import faiss
|
|
import numpy as np
|
|
import json
|
|
import re
|
|
from sklearn.feature_extraction.text import TfidfVectorizer
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
import nltk
|
|
import spacy
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
nltk.download('wordnet')
|
|
nltk.download('punkt')
|
|
|
|
try:
|
|
spacy.cli.download("en_core_web_sm")
|
|
except Exception as e:
|
|
print(f"Error downloading spacy model: {e}")
|
|
|
|
DEPLOYED = os.getenv("DEPLOYED", "false").lower() == "true"
|
|
MODEL_NAME = "tiiuae/falcon-180B-chat"
|
|
HEADERS = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
|
|
ENDPOINT_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
|
|
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."""
|
|
|
|
sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
|
|
ner_model = spacy.load("en_core_web_sm")
|
|
|
|
index = faiss.IndexFlatL2(384)
|
|
|
|
tfidf_vectorizer = TfidfVectorizer(stop_words='english')
|
|
|
|
@dataclass
|
|
class Rating:
|
|
prompt: str
|
|
response: str
|
|
ratings: list[str]
|
|
|
|
class Document:
|
|
def __init__(self, content: str, metadata: dict):
|
|
self.content = content
|
|
self.metadata = metadata
|
|
self.embedding = None
|
|
|
|
def compute_embedding(self):
|
|
self.embedding = sentence_model.encode([self.content])[0]
|
|
|
|
class DocumentStore:
|
|
def __init__(self):
|
|
self.documents = []
|
|
self.index = faiss.IndexFlatL2(384)
|
|
|
|
def add_document(self, document: Document):
|
|
document.compute_embedding()
|
|
self.documents.append(document)
|
|
self.index.add(np.array([document.embedding]))
|
|
|
|
def search(self, query: str, k: int = 5):
|
|
query_vector = sentence_model.encode([query])[0]
|
|
distances, indices = self.index.search(np.array([query_vector]), k)
|
|
return [self.documents[i] for i in indices[0]]
|
|
|
|
document_store = DocumentStore()
|
|
|
|
def extract_text_from_file(file):
|
|
if file.name.endswith('.pdf'):
|
|
reader = PyPDF2.PdfReader(file)
|
|
text = ""
|
|
for page in reader.pages:
|
|
text += page.extract_text()
|
|
elif file.name.endswith('.docx'):
|
|
doc = docx.Document(file)
|
|
text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
|
|
else:
|
|
text = file.read().decode('utf-8')
|
|
return text
|
|
|
|
def query_falcon(prompt: str, max_tokens: int = 100, temperature: float = 0.7) -> str:
|
|
payload = {
|
|
"inputs": prompt,
|
|
"parameters": {
|
|
"max_new_tokens": max_tokens,
|
|
"do_sample": True,
|
|
"temperature": temperature,
|
|
"top_p": 0.9,
|
|
"stop": ["User:"]
|
|
}
|
|
}
|
|
|
|
response = requests.post(ENDPOINT_URL, headers=HEADERS, json=payload)
|
|
if response.status_code == 200:
|
|
return response.json()[0]['generated_text']
|
|
else:
|
|
print(f"Error: {response.status_code} - {response.text}")
|
|
return "Error occurred while querying Falcon 180B."
|
|
|
|
def summarize_text(text: str, max_length: int = 150) -> str:
|
|
prompt = f"Summarize the following text in about {max_length} words:\n\n{text}\n\nSummary:"
|
|
return query_falcon(prompt, max_tokens=max_length)
|
|
|
|
def extract_legal_entities(text: str):
|
|
doc = ner_model(text)
|
|
return [ent.text for ent in doc.ents if ent.label_ in ["PERSON", "ORG", "GPE", "LAW"]]
|
|
|
|
def analyze_sentiment(text: str) -> str:
|
|
prompt = f"Analyze the sentiment of the following text. Respond with either 'Positive', 'Negative', or 'Neutral':\n\n{text}\n\nSentiment:"
|
|
return query_falcon(prompt, max_tokens=10)
|
|
|
|
def extract_keywords(text: str, top_n: int = 5):
|
|
prompt = f"Extract the top {top_n} keywords from the following text:\n\n{text}\n\nKeywords:"
|
|
response = query_falcon(prompt, max_tokens=50)
|
|
return [keyword.strip() for keyword in response.split(',')][:top_n]
|
|
|
|
def get_legal_definitions(term: str) -> str:
|
|
prompt = f"Provide a legal definition for the term '{term}':"
|
|
return query_falcon(prompt, max_tokens=100).strip()
|
|
|
|
def perform_case_law_search(query: str):
|
|
prompt = f"Perform a case law search for the following query and provide a summary of relevant cases:\n\n{query}\n\nRelevant cases:"
|
|
return query_falcon(prompt, max_tokens=200)
|
|
|
|
def generate_legal_document(document_type: str, details: dict) -> str:
|
|
prompt = f"Generate a {document_type} with the following details:\n"
|
|
for key, value in details.items():
|
|
prompt += f"{key}: {value}\n"
|
|
prompt += f"\nGenerated {document_type}:"
|
|
return query_falcon(prompt, max_tokens=500).strip()
|
|
|
|
def perform_wikipedia_search(query):
|
|
try:
|
|
search_results = wikipedia.search(query)
|
|
if search_results:
|
|
page = wikipedia.page(search_results[0])
|
|
summary = wikipedia.summary(search_results[0], sentences=3)
|
|
return f"Wikipedia: {summary}\n\nFull article: {page.url}"
|
|
else:
|
|
return "No Wikipedia results found."
|
|
except:
|
|
return "Error occurred while searching Wikipedia."
|
|
|
|
def perform_internet_search(query):
|
|
try:
|
|
search_results = list(googlesearch.search(query, num_results=3))
|
|
if search_results:
|
|
return "Internet search results:\n" + "\n".join(search_results)
|
|
else:
|
|
return "No internet search results found."
|
|
except:
|
|
return "Error occurred while performing internet search."
|
|
|
|
def chat_accordion():
|
|
with gr.Accordion("Parameters", open=False):
|
|
temperature = gr.Slider(
|
|
minimum=0.1,
|
|
maximum=1.0,
|
|
value=0.7,
|
|
step=0.1,
|
|
interactive=True,
|
|
label="Temperature",
|
|
)
|
|
top_p = gr.Slider(
|
|
minimum=0.1,
|
|
maximum=0.99,
|
|
value=0.9,
|
|
step=0.01,
|
|
interactive=True,
|
|
label="p (nucleus sampling)",
|
|
)
|
|
|
|
max_tokens = gr.Slider(
|
|
minimum=64,
|
|
maximum=1024,
|
|
value=64,
|
|
step=1,
|
|
interactive=True,
|
|
label="Max Tokens",
|
|
)
|
|
|
|
session_id = gr.Textbox(
|
|
value=uuid4,
|
|
interactive=False,
|
|
visible=False,
|
|
)
|
|
|
|
with gr.Accordion("Instructions", open=False, visible=False):
|
|
instructions = gr.Textbox(
|
|
placeholder="The Instructions",
|
|
value=DEFAULT_INSTRUCTIONS,
|
|
lines=16,
|
|
interactive=True,
|
|
label="Instructions",
|
|
max_lines=16,
|
|
show_label=False,
|
|
)
|
|
with gr.Row():
|
|
with gr.Column():
|
|
user_name = gr.Textbox(
|
|
lines=1,
|
|
label="username",
|
|
value="User",
|
|
interactive=True,
|
|
placeholder="Username: ",
|
|
show_label=False,
|
|
max_lines=1,
|
|
)
|
|
with gr.Column():
|
|
bot_name = gr.Textbox(
|
|
lines=1,
|
|
value="LexAI",
|
|
interactive=True,
|
|
placeholder="Bot Name",
|
|
show_label=False,
|
|
max_lines=1,
|
|
visible=False,
|
|
)
|
|
|
|
return temperature, top_p, instructions, user_name, bot_name, session_id, max_tokens
|
|
|
|
def format_chat_prompt(message: str, chat_history, instructions: str, user_name: str, bot_name: str):
|
|
instructions = instructions or DEFAULT_INSTRUCTIONS
|
|
instructions = instructions.strip()
|
|
prompt = instructions
|
|
for turn in chat_history:
|
|
user_message, bot_message = turn
|
|
prompt = f"{prompt}\n{user_name}: {user_message}\n{bot_name}: {bot_message}"
|
|
prompt = f"{prompt}\n{user_name}: {message}\n{bot_name}:"
|
|
return prompt
|
|
|
|
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):
|
|
if uploaded_file is not None:
|
|
document_text = extract_text_from_file(uploaded_file)
|
|
summary = summarize_text(document_text)
|
|
legal_entities = extract_legal_entities(document_text)
|
|
sentiment = analyze_sentiment(document_text)
|
|
keywords = extract_keywords(document_text)
|
|
|
|
doc = Document(content=document_text, metadata={
|
|
"filename": uploaded_file.name,
|
|
"summary": summary,
|
|
"legal_entities": legal_entities,
|
|
"sentiment": sentiment,
|
|
"keywords": keywords
|
|
})
|
|
document_store.add_document(doc)
|
|
|
|
message += f"\n[System: A document '{uploaded_file.name}' has been uploaded and processed.]"
|
|
|
|
relevant_docs = document_store.search(message)
|
|
retrieved_context = "\n".join([doc.content for doc in relevant_docs])
|
|
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
wiki_future = executor.submit(perform_wikipedia_search, message)
|
|
internet_future = executor.submit(perform_internet_search, message)
|
|
|
|
wiki_result = wiki_future.result()
|
|
internet_result = internet_future.result()
|
|
|
|
case_law_results = perform_case_law_search(message)
|
|
|
|
full_context = f"""Retrieved Documents:\n{retrieved_context}
|
|
|
|
Wikipedia Search:\n{wiki_result}
|
|
|
|
Internet Search:\n{internet_result}
|
|
|
|
Relevant Case Law:\n{case_law_results}
|
|
"""
|
|
|
|
prompt = format_chat_prompt(message, history, instructions, user_name, bot_name)
|
|
prompt += f"\nAdditional Context:\n{full_context}\n\nBased on the above information, please provide a comprehensive response:"
|
|
|
|
response = query_falcon(prompt, max_tokens=max_tokens, temperature=temperature)
|
|
|
|
response = post_process_output(response, message)
|
|
|
|
return response
|
|
|
|
def post_process_output(output: str, original_query: str) -> str:
|
|
if "define" in original_query.lower() or "meaning of" in original_query.lower():
|
|
terms = re.findall(r'\b(?!(?:the|a|an)\b)\w+', original_query)
|
|
for term in terms:
|
|
definition = get_legal_definitions(term)
|
|
output += f"\n\nLegal definition of '{term}': {definition}"
|
|
|
|
if "draft" in original_query.lower() or "create document" in original_query.lower():
|
|
doc_type_match = re.search(r'draft a (\w+)', original_query.lower())
|
|
if doc_type_match:
|
|
doc_type = doc_type_match.group(1)
|
|
details = extract_document_details(original_query)
|
|
generated_document = generate_legal_document(doc_type, details)
|
|
output += f"\n\nGenerated {doc_type.capitalize()}:\n\n{generated_document}"
|
|
|
|
if "analyze" in original_query.lower() and "document" in original_query.lower():
|
|
sentiment = analyze_sentiment(output)
|
|
output += f"\n\nOverall sentiment of the analysis: {sentiment}"
|
|
|
|
return output
|
|
|
|
def extract_document_details(query: str) -> dict:
|
|
details = {}
|
|
if "parties" in query.lower():
|
|
details["parties"] = re.search(r'parties: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
|
|
if "date" in query.lower():
|
|
details["date"] = re.search(r'date: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
|
|
if "terms" in query.lower():
|
|
details["terms"] = re.search(r'terms: (.*?)(,|\.|$)', query, re.IGNORECASE).group(1)
|
|
return details
|
|
|
|
def chat_tab():
|
|
with gr.Column():
|
|
with gr.Row():
|
|
(
|
|
temperature,
|
|
top_p,
|
|
instructions,
|
|
user_name,
|
|
bot_name,
|
|
session_id,
|
|
max_tokens
|
|
) = chat_accordion()
|
|
|
|
with gr.Column():
|
|
with gr.Blocks():
|
|
prompt_examples = [
|
|
["Analyze this contract for potential risks and summarize key points."],
|
|
["Find relevant case law for an intellectual property dispute in the software industry."],
|
|
["What are the key clauses in a non-disclosure agreement? Draft a template."],
|
|
["Predict the outcome of this employment discrimination case based on recent precedents."],
|
|
["Draft a cease and desist letter for trademark infringement. Parties: TechCorp and InnovateNow, Date: 2023-07-22"],
|
|
]
|
|
file_upload = gr.File(label="Upload Legal Document for Analysis")
|
|
gr.ChatInterface(
|
|
fn=run_chat,
|
|
chatbot=gr.Chatbot(
|
|
height=620,
|
|
render=False,
|
|
show_label=False,
|
|
avatar_images=("images/user_icon.png", "images/lexai_icon.png"),
|
|
),
|
|
textbox=gr.Textbox(
|
|
placeholder="Ask LexAI about legal matters, analyze documents, or request drafting assistance...",
|
|
render=False,
|
|
scale=7,
|
|
),
|
|
examples=prompt_examples,
|
|
additional_inputs=[
|
|
instructions,
|
|
user_name,
|
|
bot_name,
|
|
temperature,
|
|
top_p,
|
|
session_id,
|
|
max_tokens,
|
|
file_upload
|
|
],
|
|
submit_btn="Send",
|
|
stop_btn="Stop",
|
|
retry_btn="🔄 Retry",
|
|
undo_btn="↩️ Delete",
|
|
clear_btn="🗑️ Clear",
|
|
)
|
|
|
|
def introduction():
|
|
with gr.Column(scale=2):
|
|
gr.Image("images/lexai_logo.png", elem_id="banner-image", show_label=False)
|
|
with gr.Column(scale=5):
|
|
gr.Markdown(
|
|
"""# LexAI: Advanced Legal AI Assistant
|
|
**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.**
|
|
|
|
✨ This demo is powered by the state-of-the-art Falcon 180B language model and specialized legal knowledge.
|
|
|
|
🧪 LexAI offers the following key features:
|
|
1. Contract Analysis and Risk Assessment
|
|
2. Legal Research Assistant with Case Law Search
|
|
3. Predictive Litigation Analysis
|
|
4. Intelligent Legal Document Drafting
|
|
5. Legal Chatbot for Public Access
|
|
6. Document Retrieval, Analysis, and Summarization
|
|
7. Legal Entity Recognition
|
|
8. Sentiment Analysis for Legal Texts
|
|
9. Keyword Extraction from Legal Documents
|
|
10. Legal Definition Lookup
|
|
11. Wikipedia and Internet Search Integration
|
|
|
|
⚠️ **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.
|
|
"""
|
|
)
|
|
|
|
def main():
|
|
with gr.Blocks(
|
|
css="""#chat_container {height: 820px; width: 1000px; margin-left: auto; margin-right: auto;}
|
|
#chatbot {height: 600px; overflow: auto;}
|
|
#create_container {height: 750px; margin-left: 0px; margin-right: 0px;}
|
|
#tokenizer_renderer span {white-space: pre-wrap}
|
|
"""
|
|
) as demo:
|
|
with gr.Row():
|
|
introduction()
|
|
with gr.Row():
|
|
chat_tab()
|
|
|
|
return demo
|
|
|
|
def start_demo():
|
|
demo = main()
|
|
if DEPLOYED:
|
|
demo.queue().launch(show_api=False)
|
|
else:
|
|
demo.queue().launch(share=True)
|
|
|
|
if __name__ == "__main__":
|
|
start_demo() |
