Update app.py

app.py

@@ -3,12 +3,14 @@ os.system('pip install transformers')
 os.system('pip install datasets')
 os.system('pip install gradio')
 os.system('pip install minijinja')
+os.system('pip install PyMuPDF')
 
 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")
 
@@ -65,7 +67,6 @@ def generate_case_outcome(prosecutor_response, defense_response):
     return evaluation
 
 def score_argument_from_outcome(outcome, argument):
-    # Simplified scoring based on keywords in the outcome
     if "Prosecutor" in outcome:
         prosecutor_score = outcome.count("Prosecutor") * 2
         if "won" in outcome and "Prosecutor" in outcome:
@@ -167,7 +168,6 @@ footer {
 }
 """
 
-# Function to facilitate the conversation between the two chatbots
 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]
@@ -191,8 +191,35 @@ def chat_between_bots(system_message1, system_message2, max_tokens, temperature,
     
     return response1, response2, history1, history2, shared_history, outcome, prosecutor_score_color, defense_score_color
 
-def update_pdf_gallery(pdf_files):
-    return pdf_files
+def extract_text_from_pdf(pdf_file):
+    text = ""
+    doc = fitz.open(pdf_file)
+    for page in doc:
+        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
+    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"]:
@@ -202,21 +229,42 @@ def add_message(history, message):
     return history, gr.MultimodalTextbox(value=None, interactive=False)
 
 def bot(history):
-    response = "**That's cool!**"
-    history[-1][1] = ""
-    for character in response:
-        history[-1][1] += character
-        time.sleep(0.05)
+    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
 
 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([])
     pdf_files = gr.State([])
+    pdf_text = gr.State("")
     
     with gr.Tab("Argument Evaluation"):
         message = gr.Textbox(label="Case to Argue")
@@ -239,15 +287,25 @@ with gr.Blocks(css=custom_css) as demo:
         
         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_history, shared_argument, prosecutor_score_color, defense_score_color])
+        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("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_upload_btn.click(update_pdf_gallery, inputs=[pdf_upload], outputs=[pdf_gallery, pdf_files])
+        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(