Update app.py

app.py

@@ -6,7 +6,8 @@ import numpy as np
 from transformers import pipeline
 
 dataset = load_dataset("lex_glue", "scotus")
-corpus = [doc['text'] for doc in dataset['train'].select(range(200))]
+corpus_data = dataset['train'].select(range(200))
+corpus = [doc['text'] for doc in corpus_data]
 
 embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 corpus_embeddings = embedder.encode(corpus, convert_to_numpy=True)
@@ -20,7 +21,8 @@ gen_pipeline = pipeline("text2text-generation", model="facebook/bart-large-cnn")
 def rag_query(user_question):
     question_embedding = embedder.encode([user_question])
     _, indices = index.search(np.array(question_embedding), k=3)
-    context = " ".join([corpus[i] for i in indices[0]])
+    valid_indices = [i for i in indices[0] if i < len(corpus)]
+    context = " ".join([corpus[i] for i in valid_indices])
     prompt = f"Question: {user_question}\nContext: {context}\nAnswer:"
     result = gen_pipeline(prompt, max_length=250, do_sample=False)[0]['generated_text']
     return result
@@ -35,14 +37,14 @@ iface = gr.Interface(
     fn=chatbot_interface,
     inputs=[
         gr.Textbox(lines=2, placeholder="Enter your legal question here...", label="Your Question"),
-        gr.State([])  # Keeps history
+        gr.State([])  # Session state to store history
     ],
     outputs=[
         gr.Textbox(label="Chat History", lines=20, interactive=False),
         gr.State()
     ],
     title="🧑‍⚖️ Legal Assistant Chatbot",
-    description="Ask legal questions based on case data (LexGLUE - SCOTUS subset). The bot will retrieve relevant context and answer your question."
+    description="Ask legal questions based on case data (LexGLUE - SCOTUS subset). The bot retrieves context and generates an answer."
 )
 
 iface.launch()