Update app.py

app.py

@@ -4,17 +4,28 @@ os.system('pip install datasets')
 os.system('pip install gradio')
 os.system('pip install minijinja')
 os.system('pip install PyMuPDF')
+os.system('pip install pdf2image')
+os.system('pip install gradio_pdf')
 
 import gradio as gr
 from huggingface_hub import InferenceClient
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM
 from datasets import load_dataset
 import fitz  # PyMuPDF
+from pdf2image import convert_from_path
+from gradio_pdf import PDF
+from pathlib import Path
+
+dir_ = Path(__file__).parent
 
 dataset = load_dataset("ibunescu/qa_legal_dataset_train")
 
-# Use a pipeline as a high-level helper
-pipe = pipeline("fill-mask", model="nlpaueb/legal-bert-base-uncased")
+# 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")
 
@@ -36,36 +47,44 @@ def respond(
 
     messages.append({"role": "user", "content": message})
 
-    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)]
+    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 = ""
-    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
+    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_winner(outcome):
+def determine_outcome(outcome):
     prosecutor_count = outcome.split().count("Prosecutor")
     defense_count = outcome.split().count("Defense")
     if prosecutor_count > defense_count:
@@ -164,9 +183,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)
+    winner = determine_outcome(outcome)
     
-    return response1, response2, history1, history2, shared_history, outcome, winner
+    return response1, response2, history1, history2, shared_history, outcome
 
 def extract_text_from_pdf(pdf_file):
     text = ""
@@ -175,20 +194,9 @@ def extract_text_from_pdf(pdf_file):
         text += page.get_text()
     return text
 
-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 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
+def ask_about_pdf(pdf_text, question, history):
+    system_message = "You are a legal expert answering questions based on the PDF content provided."
+    response = list(respond(question, history, system_message, max_tokens=512, temperature=0.6, top_p=0.95))[-1][0]
     return response
 
 def update_pdf_gallery_and_extract_text(pdf_files):
@@ -199,7 +207,7 @@ def update_pdf_gallery_and_extract_text(pdf_files):
     return pdf_files, pdf_text
 
 def get_top_10_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. Show different cases on each click."
+    prompt = "List 10 high-profile legal cases that have received significant media attention and are currently ongoing. Just a list of case names and numbers."
     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."},
@@ -221,27 +229,10 @@ def add_message(history, message):
         history.append((message["text"], None))
     return history, gr.MultimodalTextbox(value=None, interactive=True)
 
-def bot(history):
+def bot(history, message):
     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(
-        messages,
-        max_tokens=150,
-        stream=True,
-        temperature=0.6,
-        top_p=0.95,
-    ):
-        token = message.choices[0].delta.content
-        if token is not None:
-            response += token
-        history[-1][1] = response
-        yield history
+    response = list(respond(message, history, system_message, max_tokens=150, temperature=0.6, top_p=0.95))[-1][0]
+    return response, history
 
 def print_like_dislike(x: gr.LikeData):
     print(x.index, x.value, x.liked)
@@ -268,6 +259,11 @@ def ask_about_case_outcome(shared_history, question):
             response += token
     return response
 
+def qa(question: str, doc: str) -> str:
+    img = convert_from_path(doc)[0]
+    output = pipe(img, question)
+    return sorted(output, key=lambda x: x["score"], reverse=True)[0]['answer']
+
 with gr.Blocks(css=custom_css) as demo:
     history1 = gr.State([])
     history2 = gr.State([])
@@ -301,16 +297,15 @@ 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, elem_classes=["scroll-box"])
-                winner = gr.Textbox(label="Winner", interactive=False, elem_classes=["scroll-box"])
+                outcome = gr.Textbox(label="Outcome", 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, winner])
-                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, shared_argument, winner])
+                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])
     
     with gr.Tab("PDF Management"):
@@ -324,7 +319,7 @@ with gr.Blocks(css=custom_css) as demo:
 
         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)
+        pdf_ask_btn.click(qa, inputs=[pdf_question, pdf_text], outputs=pdf_answer)
     
     with gr.Tab("Chatbot"):
         chatbot = gr.Chatbot(
@@ -336,7 +331,7 @@ with gr.Blocks(css=custom_css) as demo:
         chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)
 
         chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
-        bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
+        bot_msg = chat_msg.then(bot, inputs=[history1, chat_input], outputs=[chatbot, history1])
         bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])
 
         chatbot.like(print_like_dislike, None, None)