Update app.py

app.py

@@ -5,9 +5,15 @@ import numpy as np
 import ast
 import gradio as gr
 import faiss
+
 from sentence_transformers import SentenceTransformer
 from transformers import pipeline
 
+"""
+Legal Assistant Chatbot using LexGLUE dataset and GPT-2
+Includes session memory for last 5 Q&A and sample questions for user guidance.
+"""
+
 zip_path = "lexglue-legal-nlp-benchmark-dataset.zip"
 extract_dir = "lexglue_data"
 
@@ -36,52 +42,39 @@ generator = pipeline("text-generation", model="gpt2")
 
 history = []
 
-def simplify_legal_text(text):
-    prompt = f"Simplify the following legal text into plain English:\n\n{text}"
-    simplified_text = generator(prompt, max_new_tokens=100, do_sample=False)[0]['generated_text']
-    return simplified_text.strip()
-
-sample_questions = [
-    "Can you explain the constitutional rights of a citizen in simple terms?",
-    "What does a breach of contract mean?",
-    "How do courts determine if someone is guilty of a crime?",
-    "What is the difference between civil and criminal law?",
-    "Can you explain what 'reasonable doubt' is in a criminal trial?"
-]
-
-def legal_assistant_query(query, _):
+def legal_assistant_query(query):
     query_embedding = embedder.encode([query])
     D, I = index.search(np.array(query_embedding), k=5)
+
     retrieved_docs = [corpus[i] for i in I[0]]
     context_combined = "\n\n".join(retrieved_docs[:3])
-    max_length = 1024
-    context_combined = context_combined[:max_length]
+    context_combined = context_combined[:1024]
 
     prompt = f"Given the following legal references, answer the question:\n\n{context_combined}\n\nQuestion: {query}\nAnswer:"
     result = generator(prompt, max_new_tokens=200, do_sample=True)[0]['generated_text']
     answer = result.split("Answer:")[-1].strip()
-    simplified_answer = simplify_legal_text(answer)
 
-    history.append({"question": query, "answer": simplified_answer})
+    history.append((query, answer))
     if len(history) > 5:
         history.pop(0)
 
-    history_text = "\n\n".join([f"Q: {entry['question']}\nA: {entry['answer']}" for entry in history])
-    return simplified_answer, history_text, "\n".join(sample_questions)
+    formatted_history = "\n\n".join([f"Q: {q}\nA: {a}" for q, a in history])
+    return f"{answer}\n\n---\nRecent Q&A:\n{formatted_history}"
+
+sample_questions = [
+    "What rights does a person have under the Fourth Amendment?",
+    "Explain due process in simple terms.",
+    "What is double jeopardy?",
+    "Can the police search your car without a warrant?",
+    "What is considered a fair trial?"
+]
 
 iface = gr.Interface(
     fn=legal_assistant_query,
-    inputs=[
-        gr.Textbox(lines=2, placeholder="Ask a legal question..."),
-        gr.Button("Submit")
-    ],
-    outputs=[
-        gr.Textbox(label="Legal Response"),
-        gr.Textbox(label="Session History", lines=10),
-        gr.Textbox(label="Sample Questions", lines=6)
-    ],
+    inputs=gr.Textbox(lines=2, placeholder="Ask a legal question...", label="Your Question"),
+    outputs=gr.Textbox(label="Legal Response with History"),
     title="🧑‍⚖️ Legal Assistant Chatbot",
-    description="Ask any legal question and get context-based case references using the LexGLUE dataset. The assistant will simplify legal responses into plain English and show your last 5 questions."
+    description="Ask any legal question and get context-based case.\n\n💡 Sample Questions:\n- " + "\n- ".join(sample_questions)
 )
 
 iface.launch()