Update app.py

app.py

@@ -1,136 +1,136 @@
-# 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
+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())
+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")
+device = torch.device("cpu")
+print("Device set to use CPU")
 
 
-# embedding_model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2", device="cpu")
 
 
-# summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # -1 forces CPU usage
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # -1 forces CPU usage
 
 
-# load_dotenv()
-# api_key = os.getenv("API_KEY")
+load_dotenv()
+api_key = os.getenv("API_KEY")
 
-# genai.configure(api_key=api_key)
+genai.configure(api_key=api_key)
 
 
-# gemini_model = genai.GenerativeModel(model_name="gemini-2.0-flash")
+gemini_model = genai.GenerativeModel(model_name="gemini-2.0-flash")
 
 
-# class UserQuery(BaseModel):
-#     query: str
+class UserQuery(BaseModel):
+    query: str
 
-# class RetrievedSection(BaseModel):
-#     text: str
+class RetrievedSection(BaseModel):
+    text: str
 
-# class SummarizedResponse(BaseModel):
-#     summary: str
+class SummarizedResponse(BaseModel):
+    summary: str
 
-# class FinalLLMResponse(BaseModel):
-#     response: 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 []
+# 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)
+    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)
+    query_embedding = embedding_model.encode([query.query], convert_to_numpy=True)
 
     
-#     distances, indices = index.search(query_embedding, top_k)
+    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)
-#     ]
+    retrieved_sections = [
+        RetrievedSection(text=doc_texts[i]) for i in indices[0] if i < len(doc_texts)
+    ]
 
-#     return retrieved_sections
+    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
+# 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)
+# 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")
+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.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:")
+    st.title("🧑‍⚖️ Legal Chatbot")
+    st.markdown("### Ask your legal question below:")
 
-#     user_query = st.text_input("Enter your legal query:")
+    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...")
+    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)
+            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
+            if not retrieved_sections:
+                st.warning("No relevant legal documents found. Try refining your query.")
+                return
 
-#             summarized_sections = summarize_text(retrieved_sections)
+            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)
+            # 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}")
+            # 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("### ✨ 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)
+            st.markdown("### 🤖 AI-Enhanced Response")
+            st.text_area("Final Answer:", llm_response.response, height=150)
 
-# if __name__ == "__main__":
-#     main()
+if __name__ == "__main__":
+    main()