Update app.py

app.py

@@ -1,134 +1,220 @@
-from sentence_transformers import SentenceTransformer
-from transformers import pipeline
-from pydantic import BaseModel
-import faiss
-import numpy as np
-import streamlit as st
-from typing import List
-import os
-from dotenv import load_dotenv
-import google.generativeai as genai
-import torch
-import asyncio
-
-
-try:
-    asyncio.get_running_loop()
-except RuntimeError:
-    asyncio.set_event_loop(asyncio.new_event_loop())
-
-
-device = torch.device("cpu")
-print("Device set to use CPU")
-
-
-embedding_model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
-
-
-summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # -1 forces CPU usage
-
-
-load_dotenv()
-api_key = os.getenv("API_KEY")
-
-genai.configure(api_key=api_key)
-
-
-gemini_model = genai.GenerativeModel(model_name="gemini-2.0-flash")
-
-
-class UserQuery(BaseModel):
-    query: str
-
-class RetrievedSection(BaseModel):
-    text: str
-
-class SummarizedResponse(BaseModel):
-    summary: str
-
-class FinalLLMResponse(BaseModel):
-    response: str
-
-# Query Agent
-def query_legal_documents(query: UserQuery, top_k=3) -> List[RetrievedSection]:
-    if not os.path.exists("faiss_index.idx") or not os.path.exists("doc_texts.npy"):
-        st.error("FAISS index or document data not found.")
-        return []
-
-    
-    index = faiss.read_index("faiss_index.idx")
-    doc_texts = np.load("doc_texts.npy", allow_pickle=True)
-
-    
-    query_embedding = embedding_model.encode([query.query], convert_to_numpy=True)
-
-    
-    distances, indices = index.search(query_embedding, top_k)
-
-    
-    retrieved_sections = [
-        RetrievedSection(text=doc_texts[i]) for i in indices[0] if i < len(doc_texts)
-    ]
-
-    return retrieved_sections
-
-# Summarization Agent
-def summarize_text(text_sections: List[RetrievedSection]) -> List[SummarizedResponse]:
-    summarized_results = [
-        SummarizedResponse(
-            summary=summarizer(section.text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-        )
-        for section in text_sections
-    ]
-    return summarized_results
-
-# LLM Agent to refine response
-def generate_llm_response(summary_text: str) -> FinalLLMResponse:
-    response = gemini_model.generate_content(f"Provide a **brief** response. Do not use any special formatting like **. Here is the input:\n\n{summary_text}")
-    return FinalLLMResponse(response=response.text)
-
-
-def main():
-    st.set_page_config(page_title="Legal Chatbot", layout="wide")
-
-    
-    st.sidebar.title("Legal Chatbot Settings")
-    st.sidebar.write("This chatbot helps with legal queries by retrieving relevant legal documents, summarizing them, and generating AI-enhanced responses.")
-
-    
-    st.title("🧑‍⚖️ Legal Chatbot")
-    st.markdown("### Ask your legal question below:")
-
-    user_query = st.text_input("Enter your legal query:")
-
-    if st.button("Submit", use_container_width=True):
-        if user_query:
-            st.info("Processing your request...")
-
-            query_obj = UserQuery(query=user_query)
-            retrieved_sections = query_legal_documents(query_obj)
-
-            if not retrieved_sections:
-                st.warning("No relevant legal documents found. Try refining your query.")
-                return
-
-            summarized_sections = summarize_text(retrieved_sections)
-
-            # Combine summaries for LLM
-            combined_summary = "\n".join([res.summary for res in summarized_sections])
-            llm_response = generate_llm_response(combined_summary)
-
-            # Display results
-            st.markdown("### 📖 Retrieved Data from Knowledge Base")
-            for section in retrieved_sections:
-                st.markdown(f"🔹 {section.text}")
-
-            st.markdown("### ✨ Summarized Response")
-            for res in summarized_sections:
-                st.markdown(f"✅ {res.summary}")
-
-            st.markdown("### 🤖 AI-Enhanced Response")
-            st.text_area("Final Answer:", llm_response.response, height=150)
-
-if __name__ == "__main__":
-    main()
+# from sentence_transformers import SentenceTransformer
+# from transformers import pipeline
+# from pydantic import BaseModel
+# import faiss
+# import numpy as np
+# import streamlit as st
+# from typing import List
+# import os
+# from dotenv import load_dotenv
+# import google.generativeai as genai
+# import torch
+# import asyncio
+
+
+# try:
+#     asyncio.get_running_loop()
+# except RuntimeError:
+#     asyncio.set_event_loop(asyncio.new_event_loop())
+
+
+# device = torch.device("cpu")
+# print("Device set to use CPU")
+
+
+# embedding_model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
+
+
+# summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # -1 forces CPU usage
+
+
+# load_dotenv()
+# api_key = os.getenv("API_KEY")
+
+# genai.configure(api_key=api_key)
+
+
+# gemini_model = genai.GenerativeModel(model_name="gemini-2.0-flash")
+
+
+# class UserQuery(BaseModel):
+#     query: str
+
+# class RetrievedSection(BaseModel):
+#     text: str
+
+# class SummarizedResponse(BaseModel):
+#     summary: str
+
+# class FinalLLMResponse(BaseModel):
+#     response: str
+
+# # Query Agent
+# def query_legal_documents(query: UserQuery, top_k=3) -> List[RetrievedSection]:
+#     if not os.path.exists("faiss_index.idx") or not os.path.exists("doc_texts.npy"):
+#         st.error("FAISS index or document data not found.")
+#         return []
+
+    
+#     index = faiss.read_index("faiss_index.idx")
+#     doc_texts = np.load("doc_texts.npy", allow_pickle=True)
+
+    
+#     query_embedding = embedding_model.encode([query.query], convert_to_numpy=True)
+
+    
+#     distances, indices = index.search(query_embedding, top_k)
+
+    
+#     retrieved_sections = [
+#         RetrievedSection(text=doc_texts[i]) for i in indices[0] if i < len(doc_texts)
+#     ]
+
+#     return retrieved_sections
+
+# # Summarization Agent
+# def summarize_text(text_sections: List[RetrievedSection]) -> List[SummarizedResponse]:
+#     summarized_results = [
+#         SummarizedResponse(
+#             summary=summarizer(section.text, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
+#         )
+#         for section in text_sections
+#     ]
+#     return summarized_results
+
+# # LLM Agent to refine response
+# def generate_llm_response(summary_text: str) -> FinalLLMResponse:
+#     response = gemini_model.generate_content(f"Provide a **brief** response. Do not use any special formatting like **. Here is the input:\n\n{summary_text}")
+#     return FinalLLMResponse(response=response.text)
+
+
+# def main():
+#     st.set_page_config(page_title="Legal Chatbot", layout="wide")
+
+    
+#     st.sidebar.title("Legal Chatbot Settings")
+#     st.sidebar.write("This chatbot helps with legal queries by retrieving relevant legal documents, summarizing them, and generating AI-enhanced responses.")
+
+    
+#     st.title("🧑‍⚖️ Legal Chatbot")
+#     st.markdown("### Ask your legal question below:")
+
+#     user_query = st.text_input("Enter your legal query:")
+
+#     if st.button("Submit", use_container_width=True):
+#         if user_query:
+#             st.info("Processing your request...")
+
+#             query_obj = UserQuery(query=user_query)
+#             retrieved_sections = query_legal_documents(query_obj)
+
+#             if not retrieved_sections:
+#                 st.warning("No relevant legal documents found. Try refining your query.")
+#                 return
+
+#             summarized_sections = summarize_text(retrieved_sections)
+
+#             # Combine summaries for LLM
+#             combined_summary = "\n".join([res.summary for res in summarized_sections])
+#             llm_response = generate_llm_response(combined_summary)
+
+#             # Display results
+#             st.markdown("### 📖 Retrieved Data from Knowledge Base")
+#             for section in retrieved_sections:
+#                 st.markdown(f"🔹 {section.text}")
+
+#             st.markdown("### ✨ Summarized Response")
+#             for res in summarized_sections:
+#                 st.markdown(f"✅ {res.summary}")
+
+#             st.markdown("### 🤖 AI-Enhanced Response")
+#             st.text_area("Final Answer:", llm_response.response, height=150)
+
+# if __name__ == "__main__":
+#     main()
+
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+import faiss
+import numpy as np
+import streamlit as st
+import os
+from dotenv import load_dotenv
+import google.generativeai as genai
+import torch
+
+# Set device to CPU
+device = "cpu"
+
+# Load models once
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2-int8", device=device, normalize_embeddings=True)
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)
+
+# Load API Key
+load_dotenv()
+api_key = os.getenv("API_KEY")
+genai.configure(api_key=api_key)
+gemini_model = genai.GenerativeModel(model_name="gemini-2.0-flash")
+
+
+
+# Query Legal Documents
+def query_legal_documents(query: str, top_k=3):
+    if faiss_index is None or doc_texts is None:
+        st.error("FAISS index or document data not found.")
+        return []
+
+    query_embedding = embedding_model.encode([query])
+    distances, indices = faiss_index.search(query_embedding, top_k)
+
+    return [doc_texts[i] for i in indices[0] if i < len(doc_texts)]
+
+# Summarization Agent (Batch Processing)
+def summarize_text(text_sections):
+    texts = [section for section in text_sections]
+    summaries = summarizer(texts, max_length=100, min_length=30, do_sample=False)
+    return [summary["summary_text"] for summary in summaries]
+
+# LLM Agent (Skip if Summaries are Sufficient)
+def generate_llm_response(summary_text):
+    if len(summary_text) < 200:
+        return summary_text  # Skip LLM for short summaries
+    response = gemini_model.generate_content(summary_text)
+    return response.text
+
+# Streamlit App
+def main():
+    st.set_page_config(page_title="Legal Chatbot", layout="wide")
+    st.title("🧑‍⚖️ Legal Chatbot")
+    user_query = st.text_input("Enter your legal query:")
+
+    if st.button("Submit"):
+        if user_query:
+            st.info("Processing your request...")
+            retrieved_sections = query_legal_documents(user_query)
+
+            if not retrieved_sections:
+                st.warning("No relevant legal documents found.")
+                return
+
+            summarized_sections = summarize_text(retrieved_sections)
+            combined_summary = "\n".join(summarized_sections)
+            final_response = generate_llm_response(combined_summary)
+
+            st.markdown("### 📖 Retrieved Data")
+            for section in retrieved_sections:
+                st.markdown(f"🔹 {section}")
+
+            st.markdown("### ✨ Summarized Response")
+            for summary in summarized_sections:
+                st.markdown(f"✅ {summary}")
+
+            st.markdown("### 🤖 AI-Enhanced Response")
+            st.text_area("Final Answer:", final_response, height=150)
+
+if __name__ == "__main__":
+    main()
+
+
+
+