File size: 3,443 Bytes
6673f5d f01536b 6673f5d 42446f6 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 5e83e15 6673f5d 42446f6 6673f5d 42446f6 6673f5d 42446f6 6673f5d 42446f6 f01536b 42446f6 f01536b 42446f6 6673f5d 42446f6 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 6673f5d f01536b 6673f5d | 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 | import os
import streamlit as st
from transformers import AutoTokenizer, AutoModel, pipeline
import faiss
import pickle
import pdfplumber
from docx import Document
# Initialize variables
INDEX_FILE = "simple_rag_index.pkl"
EMBEDDING_DIM = 384 # Dimension for "all-MiniLM-L6-v2"
# Initialize FAISS index and document store
if os.path.exists(INDEX_FILE):
with open(INDEX_FILE, "rb") as f:
document_index, doc_store = pickle.load(f)
else:
document_index = faiss.IndexFlatL2(EMBEDDING_DIM)
doc_store = []
# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
# Initialize a question-answering pipeline
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased", tokenizer="distilbert-base-uncased")
# Streamlit UI
st.title("Simple RAG Chatbot ")
# Function to extract text from PDF
def extract_pdf_text(file):
with pdfplumber.open(file) as pdf:
text = ""
for page in pdf.pages:
text += page.extract_text()
return text
# Function to extract text from DOCX
def extract_docx_text(file):
doc = Document(file)
text = ""
for para in doc.paragraphs:
text += para.text + "\n"
return text
# Step 1: Upload Document
uploaded_file = st.file_uploader("Upload a document (txt, pdf, docx)", type=["txt", "pdf", "docx"])
if uploaded_file:
file_extension = uploaded_file.name.split('.')[-1].lower()
if file_extension == "txt":
document_content = uploaded_file.read().decode("utf-8")
elif file_extension == "pdf":
document_content = extract_pdf_text(uploaded_file)
elif file_extension == "docx":
document_content = extract_docx_text(uploaded_file)
else:
st.error("Unsupported file type!")
document_content = None
if document_content:
st.write("Document uploaded successfully!")
doc_store.append(document_content)
# Generate embedding for the document
tokens = tokenizer(document_content, return_tensors="pt", truncation=True, max_length=512)
embedding = model(**tokens).last_hidden_state.mean(dim=1).detach().numpy()
# Verify the embedding dimension matches FAISS index
assert embedding.shape[1] == EMBEDDING_DIM, "Embedding dimension mismatch!"
# Add embedding to FAISS index
document_index.add(embedding)
# Save the updated index and document store
with open(INDEX_FILE, "wb") as f:
pickle.dump((document_index, doc_store), f)
# Step 2: Ask a Question
user_query = st.text_input("Ask a question about the document:")
if user_query and len(doc_store) > 0:
# Generate embedding for the query
query_tokens = tokenizer(user_query, return_tensors="pt", truncation=True, max_length=512)
query_embedding = model(**query_tokens).last_hidden_state.mean(dim=1).detach().numpy()
# Retrieve the most relevant document
distances, indices = document_index.search(query_embedding, k=1)
closest_doc = doc_store[indices[0][0]]
st.write("Retrieved Document Context:", closest_doc[:500], "...") # Show a snippet of the document
# Use QA pipeline to answer the query based on the retrieved document
answer = qa_pipeline(question=user_query, context=closest_doc)
st.write("Answer:", answer["answer"])
|