Update app.py

app.py

@@ -11,15 +11,11 @@ from transformers import pipeline
 from datasets import load_dataset
 import fitz  # PyMuPDF
 
-# Loading datasets
 dataset = load_dataset("ibunescu/qa_legal_dataset_train")
 
-# Different pipelines for different tasks
-qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
-summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
-mask_filling_pipeline = pipeline("fill-mask", model="nlpaueb/legal-bert-base-uncased")
+# Use a pipeline as a high-level helper
+pipe = pipeline("fill-mask", model="nlpaueb/legal-bert-base-uncased")
 
-# Inference client for chat completion
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 def respond(
@@ -71,7 +67,7 @@ def generate_case_outcome(prosecutor_response, defense_response):
 
 def determine_winner(outcome):
     if "Prosecutor" in outcome and "Defense" in outcome:
-        if outcome.count("Prosecutor") > outcome.count("Defense"):
+        if "Prosecutor" in outcome.split().count("Prosecutor") > outcome.split().count("Defense"):
             return "Prosecutor Wins"
         else:
             return "Defense Wins"
@@ -183,8 +179,20 @@ def extract_text_from_pdf(pdf_file):
     return text
 
 def ask_about_pdf(pdf_text, question):
-    result = qa_pipeline(question=question, context=pdf_text)
-    return result['answer']
+    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 expert answering questions based on the PDF content 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
 
 def update_pdf_gallery_and_extract_text(pdf_files):
     if len(pdf_files) > 0:
@@ -194,8 +202,20 @@ def update_pdf_gallery_and_extract_text(pdf_files):
     return pdf_files, pdf_text
 
 def get_top_10_cases():
-    result = summarization_pipeline("Top 10 current legal cases in the country", max_length=150, min_length=50, do_sample=False)
-    return result[0]['summary_text']
+    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."},
+         {"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
 
 def add_message(history, message):
     for x in message["files"]:
@@ -236,8 +256,20 @@ def save_conversation(history1, history2, shared_history):
     return history1, history2, shared_history
 
 def ask_about_case_outcome(shared_history, question):
-    result = qa_pipeline(question=question, context=shared_history)
-    return result['answer']
+    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([])
@@ -283,14 +315,6 @@ with gr.Blocks(css=custom_css) as demo:
                 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])
-                
-                # Inner HTML for asking about the case outcome
-                with gr.Row():
-                    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)
     
     with gr.Tab("PDF Management"):
         pdf_upload = gr.File(label="Upload Case Files (PDF)", file_types=[".pdf"])
@@ -319,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()