Spaces:
Runtime error
Runtime error
File size: 3,765 Bytes
8f351ba 0642b48 8f351ba 882ca0c 8f351ba 882ca0c 8f351ba 882ca0c 8f351ba 29ef18c 8f351ba 882ca0c 8f351ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import os
import time
import streamlit as st
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferWindowMemory
from langchain.chains import ConversationalRetrievalChain
from langchain_together import Together
from footer import footer # Ensure this module is present in the working directory
# Set Streamlit configuration
st.set_page_config(page_title="AI Legal App", layout="centered")
# Display a logo or banner (replace with a local image or URL)
col1, col2, col3 = st.columns([1, 30, 1])
with col2:
st.image("https://github.com/Nike-one/BharatLAW/blob/master/images/banner.png?raw=true", use_column_width=True)
def hide_hamburger_menu():
st.markdown("""
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
""", unsafe_allow_html=True)
hide_hamburger_menu()
# Initialize session state
if "messages" not in st.session_state:
st.session_state.messages = []
if "memory" not in st.session_state:
st.session_state.memory = ConversationBufferWindowMemory(k=2, memory_key="chat_history", return_messages=True)
@st.cache_resource
def load_embeddings():
"""Load and cache the embeddings model."""
return HuggingFaceEmbeddings(model_name="law-ai/InLegalBERT")
embeddings = load_embeddings()
db = FAISS.load_local("ipc_embed_db", embeddings, allow_dangerous_deserialization=True)
db_retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})
prompt_template = """
<s>[INST]
As a legal chatbot specializing in Indian law, your responses must be concise and accurate:
- Provide bullet points summarizing key legal aspects.
- Avoid assumptions or overly specific advice unless requested.
- Clarify any common misconceptions.
- Keep responses aligned with general legal principles.
CONTEXT: {context}
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:
</s>[INST]
"""
prompt = PromptTemplate(template=prompt_template,
input_variables=['context', 'question', 'chat_history'])
api_key = os.getenv('TOGETHER_API_KEY')
llm = Together(model="mistralai/Mixtral-8x22B-Instruct-v0.1", temperature=0.5, max_tokens=1024, together_api_key=api_key)
qa = ConversationalRetrievalChain.from_llm(llm=llm, memory=st.session_state.memory, retriever=db_retriever, combine_docs_chain_kwargs={'prompt': prompt})
def extract_answer(full_response):
"""Extracts the assistant's answer from the response."""
return full_response.strip()
def reset_conversation():
st.session_state.messages = []
st.session_state.memory.clear()
for message in st.session_state.messages:
with st.chat_message(message["role"]):
st.write(message["content"])
input_prompt = st.chat_input("Ask your legal query...")
if input_prompt:
with st.chat_message("user"):
st.markdown(f"**You:** {input_prompt}")
st.session_state.messages.append({"role": "user", "content": input_prompt})
with st.chat_message("assistant"):
with st.spinner("Analyzing..."):
result = qa.invoke(input=input_prompt)
message_placeholder = st.empty()
answer = extract_answer(result["answer"])
# Simulated typing effect
response = ""
for char in answer:
response += char
time.sleep(0.02)
message_placeholder.markdown(response + " |", unsafe_allow_html=True)
st.session_state.messages.append({"role": "assistant", "content": answer})
if st.button('🗑️ Reset Chat', on_click=reset_conversation):
st.experimental_rerun()
footer()
|