Update app.py

app.py

@@ -9,7 +9,6 @@ import gradio as gr
 from huggingface_hub import InferenceClient
 from transformers import pipeline
 from datasets import load_dataset
-import time
 import fitz  # PyMuPDF
 
 dataset = load_dataset("ibunescu/qa_legal_dataset_train")
@@ -51,7 +50,7 @@ def respond(
         yield response, history + [(message, response)]
 
 def generate_case_outcome(prosecutor_response, defense_response):
-    prompt = f"Prosecutor's Argument: {prosecutor_response}\nDefense Attorney's Argument: {defense_response}\n\nEvaluate both arguments and determine who won the case. Provide reasons for your decision."
+    prompt = f"Prosecutor's Argument: {prosecutor_response}\nDefense Attorney's Argument: {defense_response}\n\nEvaluate both arguments, point out the strengths and weaknesses, and determine who won the case. Provide reasons for your decision."
     evaluation = ""
     for message in client.chat_completion(
         [{"role": "system", "content": "You are a legal expert evaluating the arguments presented by the prosecution and the defense."},
@@ -66,30 +65,18 @@ def generate_case_outcome(prosecutor_response, defense_response):
             evaluation += token
     return evaluation
 
-def score_argument_from_outcome(outcome, argument):
-    if "Prosecutor" in outcome:
-        prosecutor_score = outcome.count("Prosecutor") * 2
-        if "won" in outcome and "Prosecutor" in outcome:
-            prosecutor_score += 10
+def determine_winner(outcome):
+    if "Prosecutor" in outcome and "Defense" in outcome:
+        if "Prosecutor" in outcome.split().count("Prosecutor") > outcome.split().count("Defense"):
+            return "Prosecutor Wins"
+        else:
+            return "Defense Wins"
+    elif "Prosecutor" in outcome:
+        return "Prosecutor Wins"
+    elif "Defense" in outcome:
+        return "Defense Wins"
     else:
-        prosecutor_score = 0
-    
-    if "Defense" in outcome:
-        defense_score = outcome.count("Defense") * 2
-        if "won" in outcome and "Defense" in outcome:
-            defense_score += 10
-    else:
-        defense_score = 0
-
-    return prosecutor_score if "Prosecutor" in argument else defense_score
-
-def color_code(score):
-    if score > 50:
-        return "green"
-    elif score > 30:
-        return "yellow"
-    else:
-        return "red"
+        return "No clear winner"
 
 # Custom CSS for white background and black text for input and output boxes
 custom_css = """
@@ -111,6 +98,7 @@ body {
     background-color: #ffffff !important;
     border-color: #ffffff !important;
     color: #000000 !important;
+    margin: 5px;
 }
 .gr-button:hover {
     background-color: #ffffff !important;
@@ -179,17 +167,9 @@ def chat_between_bots(system_message1, system_message2, max_tokens, temperature,
     response2 = response2[:max_length]
     
     outcome = generate_case_outcome(response1, response2)
+    winner = determine_winner(outcome)
     
-    score1 = score_argument_from_outcome(outcome, "Prosecutor")
-    score2 = score_argument_from_outcome(outcome, "Defense")
-    
-    prosecutor_color = color_code(score1)
-    defense_color = color_code(score2)
-    
-    prosecutor_score_color = f"<div class='score-box' style='background-color:{prosecutor_color};'>Score: {score1}</div>"
-    defense_score_color = f"<div class='score-box' style='background-color:{defense_color};'>Score: {score2}</div>"
-    
-    return response1, response2, history1, history2, shared_history, outcome, prosecutor_score_color, defense_score_color
+    return response1, response2, history1, history2, shared_history, outcome, winner
 
 def extract_text_from_pdf(pdf_file):
     text = ""
@@ -222,7 +202,7 @@ def update_pdf_gallery_and_extract_text(pdf_files):
     return pdf_files, pdf_text
 
 def get_top_10_cases():
-    prompt = "Give me a list of 5 cases with names and cases number and a source link with any State Court of Appeals."
+    prompt = "Give me a list of the current top 10 cases in the country being discussed by the top lawyers in the country."
     response = ""
     for message in client.chat_completion(
         [{"role": "system", "content": "You are a legal expert providing information about top legal cases."},
@@ -275,6 +255,22 @@ def reset_conversation():
 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:"
+    response = ""
+    for message in client.chat_completion(
+        [{"role": "system", "content": "You are a legal expert answering questions based on the case outcome provided."},
+         {"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:
+            response += token
+    return response
+
 with gr.Blocks(css=custom_css) as demo:
     history1 = gr.State([])
     history2 = gr.State([])
@@ -308,13 +304,16 @@ with gr.Blocks(css=custom_css) as demo:
                     with gr.Column(scale=1):
                         defense_score_color = gr.HTML()
                 
-                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")
+                shared_argument = gr.Textbox(label="Case Outcome", interactive=True, elem_classes=["scroll-box"])
+                winner = gr.Textbox(label="Winner", interactive=False, 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_argument, prosecutor_score_color, defense_score_color])
-                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, shared_argument])
+                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, winner])
+                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, shared_argument, winner])
                 save_btn.click(save_conversation, inputs=[history1, history2, shared_history], outputs=[history1, history2, shared_history])
     
     with gr.Tab("PDF Management"):
@@ -344,6 +343,13 @@ with gr.Blocks(css=custom_css) as demo:
         bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
 
         chatbot.like(print_like_dislike, None, None)
+    
+    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()