Update app.py

app.py

@@ -12,182 +12,37 @@ from datasets import load_dataset
 import time
 import fitz  # PyMuPDF
 
-dir_ = Path(__file__).parent
-
 dataset = load_dataset("ibunescu/qa_legal_dataset_train")
 
-# Load the BERT model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
-model = AutoModelForMaskedLM.from_pretrained("google-bert/bert-base-uncased")
-
-# Create the fill-mask pipeline
-pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer)
-
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    try:
-        response = ""
-        for message in client.chat_completion(
-            messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = message.choices[0].delta.content
-            if token is not None:
-                response += token
-            yield response, history + [(message, response)]
-    except Exception as e:
-        print(f"Error during chat completion: {e}")
-        yield "An error occurred during the chat completion.", history
-
-def generate_case_outcome(prosecutor_response, defense_response):
-    prompt = f"Prosecutor's arguments: {prosecutor_response}\n\nDefense's arguments: {defense_response}\n\nProvide details on who won the case and why. Provide reasons for your decision and provide a link to the source of the case."
-    evaluation = ""
-    try:
-        for message in client.chat_completion(
-            [{"role": "system", "content": "You are a legal expert evaluating the details of the case presented by the prosecution and the defense."},
-             {"role": "user", "content": prompt}],
-            max_tokens=512,
-            stream=True,
-            temperature=0.6,
-            top_p=0.95,
-        ):
-            token = message.choices[0].delta.content
-            if token is not None:
-                evaluation += token
-    except Exception as e:
-        print(f"Error during case outcome generation: {e}")
-        return "An error occurred during the case outcome generation."
     return evaluation
 
-def determine_outcome(outcome):
-    prosecutor_count = outcome.split().count("Prosecutor")
-    defense_count = outcome.split().count("Defense")
-    if prosecutor_count > defense_count:
-        return "Prosecutor Wins"
-    elif defense_count > prosecutor_count:
-        return "Defense Wins"
-    else:
-        return "No clear winner"
+def score_argument_from_outcome(outcome, argument):
 
-# Custom CSS for white background and black text for input and output boxes
-custom_css = """
-body {
-    background-color: #ffffff;
-    color: #000000;
-    font-family: Arial, sans-serif;
-}
-.gradio-container {
-    max-width: 1000px;
-    margin: 0 auto;
-    padding: 20px;
-    background-color: #ffffff;
-    border: 1px solid #e0e0e0;
-    border-radius: 8px;
-    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
-}
-.gr-button {
-    background-color: #ffffff !important;
-    border-color: #ffffff !important;
-    color: #000000 !important;
-    margin: 5px;
-}
-.gr-button:hover {
-    background-color: #ffffff !important;
-    border-color: #004085 !important;
-}
-.gr-input, .gr-textbox, .gr-slider, .gr-markdown, .gr-chatbox {
-    border-radius: 4px;
-    border: 1px solid #ced4da;
-    background-color: #ffffff !important;
-    color: #000000 !important;
-}
-.gr-input:focus, .gr-textbox:focus, .gr-slider:focus {
-    border-color: #ffffff;
-    outline: 0;
-    box-shadow: 0 0 0 0.2rem rgba(255, 255, 255, 1.0);
-}
-#flagging-button {
-    display: none;
-}
-footer {
-    display: none;
-}
-.chatbox .chat-container .chat-message {
-    background-color: #ffffff !important;
-    color: #000000 !important;
-}
-.chatbox .chat-container .chat-message-input {
-    background-color: #ffffff !important;
-    color: #000000 !important;
-}
-.gr-markdown {
-    background-color: #ffffff !important;
-    color: #000000 !important;
-}
-.gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6, .gr-markdown p, .gr-markdown ul, .gr-markdown ol, .gr-markdown li {
-    color: #000000 !important;
-}
-.score-box {
-    width: 60px;
-    height: 60px;
-    display: flex;
-    align-items: center;
-    justify-content: center;
-    font-size: 12px;
-    font-weight: bold;
-    color: black;
-    margin: 5px;
-}
-.scroll-box {
-    max-height: 200px;
-    overflow-y: scroll;
-    border: 1px solid #ced4da;
-    padding: 10px;
-    border-radius: 4px;
+    if "Prosecutor" in outcome:
+        prosecutor_score = outcome.count("Prosecutor") * 2
+        if "won" in outcome and "Prosecutor" in outcome:
 }
 """
 
+
 def chat_between_bots(system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message):
     response1, history1 = list(respond(message, history1, system_message1, max_tokens, temperature, top_p))[-1]
     response2, history2 = list(respond(message, history2, system_message2, max_tokens, temperature, top_p))[-1]
-    shared_history.append(f"Prosecutor: {response1}")
-    shared_history.append(f"Defense Attorney: {response2}")
-    
-    max_length = max(len(response1), len(response2))
-    response1 = response1[:max_length]
-    response2 = response2[:max_length]
     
-    outcome = generate_case_outcome(response1, response2)
-    winner = determine_outcome(outcome)
-    
-    return response1, response2, history1, history2, shared_history, outcome
+    return response1, response2, history1, history2, shared_history, outcome, prosecutor_score_color, defense_score_color
+
+def extract_text_from_pdf(pdf_file):
+    text = ""
+    doc = fitz.open(pdf_file)
+    for page in doc:
+        text += page.get_text()
+    return text
 
-def get_top_5_cases():
-    prompt = "List 5 random high-profile legal cases that have received significant media attention and are currently ongoing. Just a list of case names and numbers."
+def ask_about_pdf(pdf_text, question):
+    prompt = f"PDF Content: {pdf_text}\n\nQuestion: {question}\n\nAnswer:"
     response = ""
     for message in client.chat_completion(
-        [{"role": "system", "content": "You are a legal research expert, able to provide information about high-profile legal cases."},
+        [{"role": "system", "content": "You are a legal expert answering questions based on the PDF content provided."},
          {"role": "user", "content": prompt}],
         max_tokens=512,
         stream=True,
@@ -199,86 +54,79 @@ def get_top_5_cases():
             response += token
     return response
 
+def update_pdf_gallery_and_extract_text(pdf_files):
+    if len(pdf_files) > 0:
+        pdf_text = extract_text_from_pdf(pdf_files[0].name)
+    else:
+        pdf_text = ""
+    return pdf_files, pdf_text
+
 def add_message(history, message):
     for x in message["files"]:
-        history.append(((x,), None))
-    if message["text"] is not None:
-        history.append((message["text"], None))
-    return history, gr.MultimodalTextbox(value=None, interactive=True)
+    return history, gr.MultimodalTextbox(value=None, interactive=False)
 
-def print_like_dislike(x: gr.LikeData):
-    print(x.index, x.value, x.liked)
-
-def reset_conversation():
-    return [], [], "", "", ""
-
-def save_conversation(history1, history2, shared_history):
-    return history1, history2, shared_history
-
-def ask_about_case_outcome(shared_history, question):
-    prompt = f"Case Outcome: {shared_history}\n\nQuestion: {question}\n\nAnswer:"
+def bot(history):
+    system_message = "You are a helpful assistant."
+    messages = [{"role": "system", "content": system_message}]
+    for val in history:
+        if val[0]:
+            messages.append({"role": "user", "content": val[0]})
+        if val[1]:
+            messages.append({"role": "assistant", "content": val[1]})
     response = ""
     for message in client.chat_completion(
-        [{"role": "system", "content": "Give the real case details of the case being argued, from a verifed source."},
-         {"role": "user", "content": prompt}],
-        max_tokens=512,
+        messages,
+        max_tokens=150,
         stream=True,
-        temperature=0.7,
+        temperature=0.6,
         top_p=0.95,
     ):
         token = message.choices[0].delta.content
         if token is not None:
             response += token
-    return response
+        history[-1][1] = response
+        yield history
+
+def print_like_dislike(x: gr.LikeData):
+    print(x.index, x.value, x.liked)
+
+def reset_conversation():
+    return [], [], "", "", ""
+
+def save_conversation(history1, history2, shared_history):
+    return history1, history2, shared_history
 
 with gr.Blocks(css=custom_css) as demo:
     history1 = gr.State([])
     history2 = gr.State([])
     shared_history = gr.State([])
-    top_10_cases = gr.State("")
-
+    pdf_files = gr.State([])
+    pdf_text = gr.State("")
+    
     with gr.Tab("Argument Evaluation"):
-        with gr.Row():
-            with gr.Column(scale=1):
-                top_5_btn = gr.Button("Give me the top 5 cases")
-                top_5_output = gr.Textbox(label="Top 5 Cases", interactive=False, elem_classes=["scroll-box"])
-                top_5_btn.click(get_top_5_cases, outputs=top_5_output)
-            with gr.Column(scale=2):
-                message = gr.Textbox(label="Case to Argue")
-                system_message1 = gr.State("You are an expert Prosecutor. Give your best arguments for the case on behalf of the prosecution.")
-                system_message2 = gr.State("You are an expert Defense Attorney. Give your best arguments for the case on behalf of the Defense.")
-                max_tokens = gr.State(512)
-                temperature = gr.State(0.5)
-                top_p = gr.State(0.95)
-                
-                with gr.Row():
-                    with gr.Column(scale=4):
-                        prosecutor_response = gr.Textbox(label="Prosecutor's Response", interactive=True, elem_classes=["scroll-box"])
-                    with gr.Column(scale=1):
-                        prosecutor_score_color = gr.HTML()
-                    
-                    with gr.Column(scale=4):
-                        defense_response = gr.Textbox(label="Defense Attorney's Response", interactive=True, elem_classes=["scroll-box"])
-                    with gr.Column(scale=1):
-                        defense_score_color = gr.HTML()
-                
-                outcome = gr.Textbox(label="Outcome", interactive=True, elem_classes=["scroll-box"])
-                
-                with gr.Row():
-                    submit_btn = gr.Button("Argue")
-                    clear_btn = gr.Button("Clear and Reset")
-                    save_btn = gr.Button("Save Conversation")
-                
-                submit_btn.click(chat_between_bots, inputs=[system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message], outputs=[prosecutor_response, defense_response, history1, history2, shared_history, outcome])
-                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, outcome])
-                save_btn.click(save_conversation, inputs=[history1, history2, shared_history], outputs=[history1, history2, shared_history])
+        message = gr.Textbox(label="Case to Argue")
+        
+        shared_argument = gr.Textbox(label="Case Outcome", interactive=True)
+        submit_btn = gr.Button("Argue")
+        clear_btn = gr.Button("Clear and Reset")
+        save_btn = gr.Button("Save Conversation")
+        
+        submit_btn.click(chat_between_bots, inputs=[system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message], outputs=[prosecutor_response, defense_response, history1, history2, shared_argument, prosecutor_score_color, defense_score_color])
+        clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, shared_argument])
+        save_btn.click(save_conversation, inputs=[history1, history2, shared_history], outputs=[history1, history2, shared_history])
     
-    with gr.Tab("Case Outcome Chat"):
-        case_question = gr.Textbox(label="Ask a Question about the Case Outcome")
-        case_answer = gr.Textbox(label="Answer", interactive=False, elem_classes=["scroll-box"])
-        ask_case_btn = gr.Button("Ask")
-
-        ask_case_btn.click(ask_about_case_outcome, inputs=[shared_history, case_question], outputs=case_answer)
-
-demo.queue()
-demo.launch()
+    with gr.Tab("PDF Management"):
+        pdf_upload = gr.File(label="Upload Case Files (PDF)", file_types=[".pdf"])
+        pdf_gallery = gr.Gallery(label="PDF Gallery")
+        pdf_view = gr.Textbox(label="PDF Content", interactive=False, elem_classes=["scroll-box"])
+        pdf_question = gr.Textbox(label="Ask a Question about the PDF")
+        pdf_answer = gr.Textbox(label="Answer", interactive=False, elem_classes=["scroll-box"])
+        pdf_upload_btn = gr.Button("Update PDF Gallery")
+        pdf_ask_btn = gr.Button("Ask")
+
+        pdf_upload_btn.click(update_pdf_gallery_and_extract_text, inputs=[pdf_upload], outputs=[pdf_gallery, pdf_text])
+        pdf_text.change(fn=lambda x: x, inputs=pdf_text, outputs=pdf_view)
+        pdf_ask_btn.click(ask_about_pdf, inputs=[pdf_text, pdf_question], outputs=pdf_answer)
+    
+    with gr.Tab("Chatbot"):
+        chatbot = gr.Chatbot(