Upload 3 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+NJ_opinions_modernbert_splitter.jsonl filter=lfs diff=lfs merge=lfs -text

NJ_opinions_modernbert_splitter.jsonl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e4d9d8c12e6784a66dfb819ddc473b76420e5d198337cd10ac0e68b93ef3466
+size 670176621

app.py

@@ -7,7 +7,7 @@ import time
 import torch
 # from retrieval import *
 import os
-from transformers import AutoTokenizer, AutoModel, pipeline , AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModel, pipeline , AutoModelForSequenceClassification, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
@@ -15,25 +15,483 @@ import pandas as pd
 import torch.nn.functional as F
 from datasets import concatenate_datasets, load_dataset, load_from_disk
 from huggingface_hub import hf_hub_download
-import random
 from contextual import ContextualAI
 from openai import AzureOpenAI
-from transformers import pipeline, set_seed
+from datetime import datetime
 
+"""
+# to switch:
+device to cuda
+enable bfloat16
 
-generator = pipeline('text-generation', model='gpt2', device_map="auto")
+"""
 
-@spaces.GPU
-def respond(message, chat_history):
-    bot_message = generator(message, max_length=10)[0]["generated_text"]
-    chat_history.append({"role": "user", "content": message})
-    chat_history.append({"role": "assistant", "content": bot_message})
-    return "", chat_history
 
-with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(type="messages")
-    msg = gr.Textbox()
-    clear = gr.ClearButton([msg, chatbot])
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
-demo.launch()
+sandbox_api_key=os.getenv('AI_SANDBOX_KEY')
+sandbox_endpoint="https://api-ai-sandbox.princeton.edu/"
+sandbox_api_version="2024-02-01"
+
+def text_prompt_call(model_to_be_used, system_prompt, user_prompt ):
+    client = AzureOpenAI(
+        api_key=sandbox_api_key,
+        azure_endpoint = sandbox_endpoint,   
+        api_version=sandbox_api_version # current api version not in preview
+    )
+    response = client.chat.completions.create(
+        model=model_to_be_used,
+        temperature=0.7, # temperature = how creative/random the model is in generating response - 0 to 1 with 1 being most creative
+        max_tokens=1000, # max_tokens = token limit on context to send to the model
+        messages=[
+            {"role": "system", "content": system_prompt}, # describes model identity and purpose
+            {"role": "user", "content": user_prompt}, # user prompt
+        ]
+    )
+    return response.choices[0].message.content
+
+
+
+api_key = os.getenv("contextual_apikey")
+base_url = "https://api.contextual.ai/v1"
+rerank_api_endpoint = f"{base_url}/rerank"
+reranker = "ctxl-rerank-en-v1-instruct"
+
+
+
+
+#client = ContextualAI (api_key = api_key, base_url = base_url)
+
+"""
+
+documents = [i["text"] for i in results]
+query = results[0]["query"]
+metadata = [i["meta_data"] for i in results]
+model = "ctxl-rerank-en-v1-instruct"
+
+query = "Find me NJ App Division Court opinions on whether officers can always order a passenger out of a car?"
+
+#instruction = "Prioritize internal sales documents over market analysis reports. More recent documents should be weighted higher. Enterprise portal content supersedes distributor communications."
+
+rerank_response = client.rerank.create(
+    query = query,
+    instruction = instruction,
+    documents = documents,
+    metadata = metadata,
+    model = model
+)
+
+print(rerank_response.to_dict())
+"""
+
+#instruction_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-7B-Instruct", torch_dtype=torch.bfloat16, device_map="auto")
+#instruction_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
+
+def update_instruction(query):
+    system_prompt_instructions = """You are given a query and an instruction. Modify the instruction to prioritize the types of documents the query specifies. If the query asks for specific details (e.g., court level, timeframe, citation importance), incorporate those details into the instruction while maintaining its original structure. If the query does not specify particular document preferences, return "not applicable."
+
+Example 1
+
+Query: Find me older appellate court opinions on whether officers can always order passengers out of a car.
+Instruction: Prioritize older appellate court opinions
+
+Example 2
+
+Query: Show me recent Supreme Court rulings on digital privacy rights.
+Output: Prioritize recent Supreme Court opinions. 
+
+Example 3
+
+Query: Find legal opinions on self-defense laws.
+Output: not applicable
+
+Example 4
+Query: Locate federal district court rulings from the last five years on employer vaccine mandates.
+Output: Prioritize federal district court rulings from the last five years.
+
+Example 5
+
+Query: Show me influential appellate court decisions on contract interpretation.
+Output: Prioritize influential appellate court decisions.
+
+Example 6
+
+Query: Find state supreme court cases that discuss the necessity of search warrants for vehicle searches.
+Output: Prioritize state supreme court cases on search warrants for vehicle searches.
+
+Example 7
+
+Query: Show me legal opinions about landlord-tenant disputes.
+Output: not applicable
+"""
+
+
+    """
+    messages = [{"role": "system", "content": system_prompt_instructions}]
+    messages.append({"role": "user", "content": "Query: " + query})
+    example = instruction_tokenizer.apply_chat_template(messages, add_generation_prompt = True, tokenize=True,pad_to_multiple_of=8, do_pan_and_scan=True, return_tensors="pt")
+    out = instruction_model.generate(example, max_new_tokens=50)
+    updated = instruction_tokenizer.decode(out[0])
+    updated = updated.split("<|im_start|>assistant")[-1].split("<|im_end|>")[0].strip()
+    """
+
+    updated = text_prompt_call("gpt-4o", system_prompt_instructions, query)
+
+    print ("UPDATED INSTRUCTION HERE", updated)
+    if updated == "not applicable":
+        return "Prioritize Supreme Court opinions or opinions from higher courts. More recent, highly cited and published documents should also be weighted higher."
+
+    return updated
+    # oh god
+
+
+
+
+
+
+def rerank_with_contextual_AI(results):
+    instruction = "Prioritize Supreme Court opinions or opinions from higher courts. More recent, highly cited and published documents should also be weighted higher."
+    #instruction = rerank_instruction
+    query = results[0]["query"]
+    docs = [i["text"] for i in results]
+    metadata = [i["meta_data"] for i in results]
+
+    # rewrite instruction if applicable
+    instruction = update_instruction(query)
+
+    rerank_response = client.rerank.create(
+        query = query,
+        instruction = instruction,
+        documents = docs,
+        metadata = metadata,
+        model = reranker
+    ).to_dict()
+    print (rerank_response)
+    # {'results': [{'index': 3, 'relevance_score': 0.39700255}, {'index': 2, 'relevance_score': 0.38903061}, {'index': 10, 'relevance_score': 0.36989796}, {'index': 8, 'relevance_score': 0.36830357}, {'index': 1, 'relevance_score': 0.36415816}, {'index': 11, 'relevance_score': 0.35778061}, {'index': 0, 'relevance_score': 0.35586735}, {'index': 4, 'relevance_score': 0.32589286}, {'index': 12, 'relevance_score': 0.32589286}, {'index': 7, 'relevance_score': 0.30931122}, {'index': 9, 'relevance_score': 0.30739796}, {'index': 13, 'relevance_score': 0.29145408}, {'index': 5, 'relevance_score': 0.2755102}, {'index': 6, 'relevance_score': 0.27295918}]}
+
+    #ok, what next?
+    reranked_docs = []
+    for i in rerank_response["results"]:
+        reranked_docs.append(results[i["index"]])
+        reranked_docs[-1]["relevance_score"] = i["relevance_score"]
+    return reranked_docs
+
+def format_metadata_for_reranking(metadata):
+    try:
+        out =  metadata["case_name"] + ", " + metadata["court_short_name"] + ", " + "year: " + metadata["date_filed"] + " citation count: " + str(metadata["citation_count"]) + ", precedential status " + metadata["precedential_status"]
+    except:
+        out = ""
+    return out
+
+
+def format_metadata_as_str(metadata):
+    try:
+        out =  metadata["case_name"] + ", " + metadata["court_short_name"] + ", " + metadata["date_filed"] + ", precedential status " + metadata["precedential_status"]
+    except:
+        out = ""
+    return out
+
+def show_user_query(user_message, history):
+    '''
+    Displays user query in the chatbot and removes from textbox.
+    :param user_message: user query inputted.
+    :param history: 2D array representing chatbot-user conversation.
+    :return:
+    '''
+    return "", history + [[user_message, None]]
+
+
+
+def run_extractive_qa(query, contexts):
+    extracted_passages = extractive_qa([{"question": query, "context": context} for context in contexts])        
+    return extracted_passages
+
+
+@spaces.GPU(duration=15)
+def respond_user_query(history):
+    '''
+    Overwrite the value of current pairing's history with generated text
+    and displays response character-by-character with some lag.
+    :param history: 2D array of chatbot history filled with user-bot interactions
+    :return: history updated with bot's latest message.
+    '''
+    start_time_global = time.time()
+
+    query = history[0][0]
+    start_time_global = time.time()
+
+    responses = run_retrieval(query)
+    print("--- run retrieval: %s seconds ---" % (time.time() - start_time_global))
+    #print (responses)
+
+    contexts = [individual_response["text"] for individual_response in responses][:NUM_RESULTS]
+    extracted_passages = run_extractive_qa(query, contexts)
+
+    for individual_response, extracted_passage in zip(responses, extracted_passages):
+        start, end = extracted_passage["start"], extracted_passage["end"]
+        # highlight text
+        text = individual_response["text"]
+        text = text[:start] + " **" + text[start:end] + "** " + text[end:]
+
+        # display queries in interface
+        formatted_response = "##### "
+        if individual_response["meta_data"]:
+            formatted_response += individual_response["meta_data"]
+        else:
+            formatted_response += individual_response["opinion_idx"]
+        formatted_response += "\n" + text + "\n\n"
+        history = history + [[None, formatted_response]]
+    print("--- Extractive QA: %s seconds ---" % (time.time() - start_time_global))
+
+    return [history, responses]
+
+def switch_to_reviewing_framework():
+    '''
+    Replaces textbox for entering user query with annotator review select.
+    :return: updated visibility for textbox and radio button props.
+    '''
+    return gr.Textbox(visible=False), gr.Dataset(visible=False), gr.Textbox(visible=True, interactive=True), gr.Button(visible=True)
+
+def reset_interface():
+    '''
+    Resets chatbot interface to original position where chatbot history,
+     reviewing is invisbile is empty and user input textbox is visible.
+    :return: textbox visibility, review radio button invisibility,
+    next_button invisibility, empty chatbot
+    '''
+
+    # remove tmp highlighted word documents
+    #for fn in os.listdir("tmp-docs"):
+    #    os.remove(os.path.join("tmp-docs", fn))
+    return gr.Textbox(visible=True), gr.Button(visible=False), gr.Textbox(visible=False, value=""), None, gr.JSON(visible=False, value=[]), gr.Dataset(visible=True)
+
+###################################################
+def mark_like(response_json, like_data: gr.LikeData):
+    index_of_msg_reviewed = like_data.index[0] - 1  # 0-indexing
+    # add liked information to res
+    response_json[index_of_msg_reviewed]["is_msg_liked"] = like_data.liked
+    return response_json
+
+"""
+def save_json(name: str, greetings: str) -> None:
+
+"""
+def register_review(history, additional_feedback, response_json):
+    '''
+    Writes user review to output file.
+    :param history: 2D array representing bot-user conversation so far.
+    :return: None, writes to output file.
+    ''' 
+
+    res = { "user_query": history[0][0], 
+        "responses": response_json,
+        "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S'), 
+        "additional_feedback": additional_feedback
+    }
+    print (res)
+
+
+# load search functionality here
+
+
+def load_bm25():
+    stemmer = Stemmer.Stemmer("english")
+    retriever = bm25s.BM25.load("NJ_index_LLM_chunking", mmap=False)
+    return retriever, stemmer # titles
+
+def run_bm25(query):
+    query_tokens = bm25s.tokenize(query, stemmer=stemmer)
+    results, scores = retriever.retrieve(query_tokens, k=5)
+    return results[0]
+
+def load_faiss_index(embeddings):
+    nb, d = embeddings.shape # database size, dimension
+    faiss_index = faiss.IndexFlatL2(d)   # build the index
+    faiss_index.add(embeddings) # add vectors to the index
+    return faiss_index
+
+#@spaces.GPU(duration=10)
+def run_dense_retrieval(query):
+    if "NV" in model_name:
+        query_prefix = "Instruct: Given a question, retrieve passages that answer the question\nQuery: "
+        max_length = 32768
+        print (query)
+        with torch.no_grad():
+            query_embeddings = model.encode([query], instruction=query_prefix, max_length=max_length)
+            query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
+        query_embeddings = query_embeddings.cpu().numpy()
+    return query_embeddings
+
+
+def load_NJ_caselaw():
+    if os.path.exists("/scratch/gpfs/ds8100/datasets/NJ_opinions_modernbert_splitter.jsonl"):
+        df = pd.read_json("/scratch/gpfs/ds8100/datasets/NJ_opinions_modernbert_splitter.jsonl", lines=True)
+    else:
+        df = pd.read_json("NJ_opinions_modernbert_splitter.jsonl", lines=True)
+    titles, chunks = [],[]
+
+    for i, row in df.iterrows():
+        texts = [i for i in row["texts"] if len(i.split()) > 25 and len(i.split()) < 750]
+        texts = [" ".join(i.strip().split()) for i in texts]
+        chunks.extend(texts)
+        titles.extend([row["id"]] * len(texts))
+    ids = list(range(len(titles)))
+    assert len(ids) == len(titles) == len(chunks)
+    return ids, titles, chunks
+
+
+def run_retrieval(query):
+    query = " ".join(query.split())
+
+    print ("query", query)
+    """
+    indices_bm25 = run_bm25(query)
+    scores_embeddings, indices_embeddings = run_dense_retrieval(query)
+    indices = list(set(indices_bm25).union(indices_embeddings))
+    #docs = [{"id":i, "text":chunks[i]} for i in indices]
+    docs = [chunks[i] for i in indices]
+    results_reranking = rerank(query, docs, indices) #results = [{"doc":docs[i], "score":probs[i], "id":indices[i]} for i in argsort]
+    """
+    start_time = time.time()
+    query_embeddings = run_dense_retrieval(query)
+    print("--- Nvidia Embedding: %s seconds ---" % (time.time() - start_time))
+    D, I = faiss_index.search(query_embeddings, 45)
+    print("--- Faiss retrieval: %s seconds ---" % (time.time() - start_time))
+
+    scores_embeddings = D[0]
+    indices_embeddings = I[0]
+
+    docs = [chunks[i] for i in indices_embeddings]
+    results = [{"id":i, "score":j} for i,j in zip(indices_embeddings, scores_embeddings)]
+
+    out_dict = []
+    covered = set()
+    for item in results:
+        tmp = {}
+        index = item["id"]
+        tmp["query"] = query
+        tmp["index"] = index #indices[index]
+        tmp["NV_score"] = item["score"]
+        tmp["opinion_idx"] = str(titles[index])
+        # and htis helps
+        if tmp["opinion_idx"] in covered:
+            continue
+        covered.add(tmp["opinion_idx"])
+        if tmp["opinion_idx"] in metadata:
+            tmp["meta_data"] = format_metadata_for_reranking(metadata[tmp["opinion_idx"]])
+        else:
+            tmp["meta_data"] = ""        
+        # so far so good
+        tmp["text"] = chunks[tmp["index"]]
+        out_dict.append(tmp)
+    print (out_dict)
+    # and now, rerank
+    #out_dict = out_dict[:NUM_RESULTS]
+    #out_dict = rerank_with_contextual_AI(out_dict)
+    return out_dict
+
+
+NUM_RESULTS = 5
+model_name = 'nvidia/NV-Embed-v2' 
+
+device = torch.device("cuda")
+#device = torch.device("cpu")
+#device = torch.device("mps")
+
+extractive_qa = pipeline("question-answering", model="ai-law-society-lab/extractive-qa-model", tokenizer="FacebookAI/roberta-large", device_map="auto", token=os.getenv('hf_token'))
+ids, titles, chunks = load_NJ_caselaw()
+
+ds = load_dataset("ai-law-society-lab/NJ_embeddings", token=os.getenv('hf_token'))["train"]
+ds = ds.with_format("np")
+print (ds)
+faiss_index = load_faiss_index(ds["embeddings"])
+
+with open("NJ_caselaw_metadata.json") as f:
+    metadata = json.load(f)
+
+ 
+
+def load_embeddings_model(model_name = "intfloat/e5-large-v2"):
+    # let's try to not rely on sentencetransformer at some point in the future?
+    if "NV" in model_name:
+        model = AutoModel.from_pretrained('nvidia/NV-Embed-v2', trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
+        #model = AutoModel.from_pretrained('nvidia/NV-Embed-v2', trust_remote_code=True, torch_dtype=torch.float16, device_map="auto")
+        model.eval()
+        #model.to(device)
+        return model
+
+if "NV" in model_name:
+    model = load_embeddings_model(model_name=model_name)
+
+
+examples = ["Can officers always order a passenger out of a car?","Find me briefs about credential searches", "Can police search an impounded car without a warrant?", "State is arguing State v. Carty is not good law"]
+
+css = """
+.svelte-i3tvor {visibility: hidden}
+.row.svelte-hrj4a0.unequal-height {
+    align-items: stretch !important
+}
+"""
+
+"""
+with gr.Blocks(css=css, theme = gr.themes.Monochrome(primary_hue="pink",)) as demo:
+    chatbot = gr.Chatbot(height="45vw", autoscroll=False)
+    query_textbox = gr.Textbox()
+    examples = gr.Examples(examples, query_textbox)
+    response_json = gr.JSON(visible=False, value=[])
+    chatbot.like(mark_like, response_json, response_json) 
+
+    with gr.Row(visible=False, equal_height=True) as downloadable_file_row:
+        @gr.render(inputs=response_json)
+        def set_download_files(response_json):
+                seen_files = set()
+                for response in response_json:
+                    if response["path_file"] not in seen_files:
+                        # use metadata title if it exists
+                        title = response["title"]
+                        if response["meta_data"]:
+                            title = response["meta_data"]
+                        gr.DownloadButton(label="Download: " + title, value = response["path_file"], visible=True, size="sm")
+                        seen_files.add(response["path_file"])
+    
+    feedback_textbox = gr.Textbox(label="Additional feedback?", visible=False)
+    next_button = gr.Button(value="Submit Feedback", visible=False)
+
+    # Handle annotator query
+    query_textbox.submit(show_user_query, [query_textbox, chatbot], [query_textbox, chatbot], queue=False).then(
+        respond_user_query, chatbot, [chatbot, response_json]).then(
+        switch_to_reviewing_framework, None, [query_textbox, examples.dataset, downloadable_file_row, feedback_textbox, next_button]
+    ) 
+
+    # Handle page reset and review save in database
+    next_button.click(register_review, [chatbot, feedback_textbox, response_json], None).then(
+        reset_interface, None, [query_textbox, next_button, feedback_textbox, chatbot, downloadable_file_row, response_json, examples.dataset])
+    #next_button.click(register_review, [chatbot, feedback_textbox, response_json], None).then(
+        reset_interface, None, [query_textbox, next_button, feedback_textbox, chatbot, downloadable_file_row, response_json, examples.dataset])
+
+
+# Launch application
+demo.launch() 
+"""
+with gr.Blocks(css=css, theme = gr.themes.Monochrome(primary_hue="pink",)) as demo:
+    chatbot = gr.Chatbot(height="45vw", autoscroll=False)
+    query_textbox = gr.Textbox()
+    #rerank_instruction = gr.Textbox(label="Rerank Instruction Prompt", value="If not otherwise specified in the query, prioritize Supreme Court opinions or opinions from higher courts. More recent, highly cited and published documents should also be weighted higher, unless otherwise specified in the query.")
+    examples = gr.Examples(examples, query_textbox)
+    response_json = gr.JSON(visible=False, value=[])
+    print (response_json)
+    chatbot.like(mark_like, response_json, response_json) 
+
+    feedback_textbox = gr.Textbox(label="Additional feedback?", visible=False)
+    next_button = gr.Button(value="Submit Feedback", visible=False)
+
+    query_textbox.submit(show_user_query, [query_textbox, chatbot], [query_textbox, chatbot], queue=False).then(
+        respond_user_query, chatbot, [chatbot, response_json]).then(
+        switch_to_reviewing_framework, None, [query_textbox, examples.dataset, feedback_textbox, next_button]
+    ) 
+
+    # Handle page reset and review save in database
+    next_button.click(register_review, [chatbot, feedback_textbox, response_json], None).then(
+        reset_interface, None, [query_textbox, next_button, feedback_textbox, chatbot, response_json, examples.dataset])
+
+# Launch application
+demo.launch()