File size: 3,768 Bytes
8f252e0 |
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 |
import streamlit as st
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.schema import Document
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from rank_bm25 import BM25Okapi
from langchain.retrievers import ContextualCompressionRetriever, BM25Retriever, EnsembleRetriever
from langchain.retrievers.document_compressors import FlashrankRerank
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_groq import ChatGroq
from langchain.prompts import ChatPromptTemplate
import hashlib
from typing import List
# Contextual Retrieval Class
class ContextualRetrieval:
def __init__(self):
self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
model_name = "BAAI/bge-large-en-v1.5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': False}
self.embeddings = HuggingFaceEmbeddings(model_name=model_name, model_kwargs=model_kwargs, encode_kwargs=encode_kwargs)
self.llm = ChatGroq(model="llama-3.2-3b-preview", temperature=0)
def process_document(self, document: str) -> List[Document]:
return self.text_splitter.create_documents([document])
def create_vectorstore(self, chunks: List[Document]) -> FAISS:
return FAISS.from_documents(chunks, self.embeddings)
def create_bm25_retriever(self, chunks: List[Document]) -> BM25Retriever:
return BM25Retriever.from_documents(chunks)
def generate_answer(self, query: str, docs: List[Document]) -> str:
prompt = ChatPromptTemplate.from_template("""
Question: {query}
Relevant Information: {chunks}
Answer:""")
messages = prompt.format_messages(query=query, chunks="\n\n".join([doc.page_content for doc in docs]))
response = self.llm.invoke(messages)
return response.content
# Streamlit UI
def main():
st.title("Interactive Document Retrieval Analysis")
st.write("Upload a document, experiment with retrieval methods, and analyze content interactively.")
# Document Upload
uploaded_file = st.file_uploader("Upload a Text Document", type=['txt', 'md'])
if uploaded_file:
document = uploaded_file.read().decode("utf-8")
st.success("Document successfully uploaded!")
# Initialize Retrieval System
cr = ContextualRetrieval()
chunks = cr.process_document(document)
vectorstore = cr.create_vectorstore(chunks)
bm25_retriever = cr.create_bm25_retriever(chunks)
# Query Input
query = st.text_input("Enter your question about the document:")
if query:
# Retrieve Results
with st.spinner("Fetching results..."):
vector_results = vectorstore.similarity_search(query, k=3)
bm25_results = bm25_retriever.get_relevant_documents(query)
vector_answer = cr.generate_answer(query, vector_results)
bm25_answer = cr.generate_answer(query, bm25_results)
# Display Results
st.subheader("Results from Vector Search")
st.write(vector_answer)
st.subheader("Results from BM25 Search")
st.write(bm25_answer)
# Display Sources
st.subheader("Top Retrieved Chunks")
st.write("**Vector Search Results:**")
for i, doc in enumerate(vector_results, 1):
st.write(f"{i}. {doc.page_content[:300]}...")
st.write("**BM25 Search Results:**")
for i, doc in enumerate(bm25_results, 1):
st.write(f"{i}. {doc.page_content[:300]}...")
# Run the Streamlit App
if __name__ == "__main__":
main() |