Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PyPDF2 import PdfReader | |
| from langchain.vectorstores.cassandra import Cassandra | |
| from langchain.indexes.vectorstore import VectorStoreIndexWrapper | |
| from langchain.embeddings import HuggingFaceEmbeddings | |
| from langchain.llms import HuggingFaceHub | |
| from langchain.text_splitter import CharacterTextSplitter | |
| import cassio | |
| from dotenv import load_dotenv | |
| import os | |
| load_dotenv() | |
| ASTRA_DB_APPLICATION_TOKEN = os.getenv("ASTRA_DB_APPLICATION_TOKEN") | |
| ASTRA_DB_ID = os.getenv("ASTRA_DB_ID") | |
| HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN") | |
| # === Streamlit UI Setup === | |
| st.set_page_config(page_title="Query PDF with Free Hugging Face Models", layout="wide") | |
| st.title("📄💬 Query PDF using LangChain + AstraDB (Free Hugging Face Models)") | |
| # === File Upload === | |
| uploaded_file = st.file_uploader("Upload your PDF", type=["pdf"]) | |
| if uploaded_file: | |
| st.success("✅ PDF uploaded successfully!") | |
| process_button = st.button("🔄 Process PDF") | |
| if process_button: | |
| # Initialize AstraDB | |
| cassio.init(token=ASTRA_DB_APPLICATION_TOKEN, database_id=ASTRA_DB_ID) | |
| # Read PDF contents | |
| pdf_reader = PdfReader(uploaded_file) | |
| raw_text = "" | |
| for page in pdf_reader.pages: | |
| content = page.extract_text() | |
| if content: | |
| raw_text += content | |
| # Split text into chunks | |
| text_splitter = CharacterTextSplitter( | |
| separator="\n", chunk_size=800, chunk_overlap=200, length_function=len | |
| ) | |
| texts = text_splitter.split_text(raw_text) | |
| # === Embeddings === | |
| embedding = HuggingFaceEmbeddings( | |
| model_name="sentence-transformers/all-MiniLM-L6-v2" | |
| ) | |
| # === Hugging Face LLM === | |
| llm = HuggingFaceHub( | |
| repo_id="mistralai/Mistral-7B-Instruct-v0.1", | |
| huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN, | |
| model_kwargs={"temperature": 0.5, "max_new_tokens": 512} | |
| ) | |
| # === Create vector store and index === | |
| vector_store = Cassandra( | |
| embedding=embedding, | |
| table_name="qa_mini_demo", | |
| session=None, | |
| keyspace=None, | |
| ) | |
| vector_store.add_texts(texts[:50]) | |
| st.success(f"📚 {len(texts[:50])} chunks embedded and stored in AstraDB.") | |
| astra_vector_index = VectorStoreIndexWrapper(vectorstore=vector_store) | |
| # === Ask Questions === | |
| st.header("🤖 Ask a question about your PDF") | |
| user_question = st.text_input("💬 Type your question here") | |
| if user_question: | |
| with st.spinner("🧠 Thinking..."): | |
| try: | |
| # Optional: show what documents are retrieved before sending to LLM | |
| retrieved_docs = vector_store.similarity_search(user_question, k=4) | |
| if not retrieved_docs: | |
| st.warning("⚠️ No relevant text chunks found for this question. Try a different question.") | |
| else: | |
| st.markdown("### 🔍 Top Relevant Chunks (raw):") | |
| for i, doc in enumerate(retrieved_docs, 1): | |
| st.code(doc.page_content[:300], language="markdown") | |
| answer = astra_vector_index.query(user_question, llm=llm) | |
| if answer.strip(): | |
| st.markdown("### 🧠 Answer:") | |
| st.write(answer.strip()) | |
| else: | |
| st.warning("⚠️ The model returned an empty response. Try rephrasing the question or check your model/API key.") | |
| except Exception as e: | |
| st.error(f"🚨 Error while generating response:\n\n{str(e)}") | |