app2.py

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()