Update app.py

app.py

@@ -13,8 +13,10 @@ import fitz  # PyMuPDF
 
 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")
+# 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")
 
 client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
@@ -179,20 +181,8 @@ def extract_text_from_pdf(pdf_file):
     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
+    result = qa_pipeline(question=question, context=pdf_text)
+    return result['answer']
 
 def update_pdf_gallery_and_extract_text(pdf_files):
     if len(pdf_files) > 0:
@@ -202,20 +192,8 @@ 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 10 random cases with case numbers and names 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
+    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']
 
 def add_message(history, message):
     for x in message["files"]:
@@ -256,20 +234,8 @@ 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
+    result = qa_pipeline(question=question, context=shared_history)
+    return result['answer']
 
 with gr.Blocks(css=custom_css) as demo:
     history1 = gr.State([])