Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,152 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
 import streamlit as st
+import PyPDF2
+import torch
+import weaviate
+from transformers import AutoTokenizer, AutoModel
+from weaviate.classes.init import Auth
+import cohere
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Load credentials from environment variables or hardcoded (replace with env vars in prod)
+WEAVIATE_URL = "vgwhgmrlqrqqgnlb1avjaa.c0.us-west3.gcp.weaviate.cloud"
+WEAVIATE_API_KEY = "7VoeYTjkOS4aHINuhllGpH4JPgE2QquFmSMn"
+COHERE_API_KEY = "LEvCVeZkqZMW1aLYjxDqlstCzWi4Cvlt9PiysqT8"
+
+# Connect to Weaviate
+client = weaviate.connect_to_weaviate_cloud(
+    cluster_url=WEAVIATE_URL,
+    auth_credentials=Auth.api_key(WEAVIATE_API_KEY),
+    headers={"X-Cohere-Api-Key": COHERE_API_KEY}
+)
+
+cohere_client = cohere.Client(COHERE_API_KEY)
+
+# Load sentence-transformer model
+tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+
+def load_pdf(file):
+    """Extract text from PDF file."""
+    reader = PyPDF2.PdfReader(file)
+    text = ''.join([page.extract_text() for page in reader.pages if page.extract_text()])
+    return text
+
+def get_embeddings(text):
+    """Generate mean pooled embedding for the input text."""
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        embeddings = model(**inputs).last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
+    return embeddings
+
+def upload_document_chunks(chunks):
+    """Insert document chunks into Weaviate collection with embeddings."""
+    doc_collection = client.collections.get("Document")
+    for chunk in chunks:
+        embedding = get_embeddings(chunk)
+        doc_collection.data.insert(
+            properties={"content": chunk},
+            vector=embedding.tolist()
+        )
+
+def query_answer(query):
+    """Search for top relevant document chunks based on query embedding."""
+    query_embedding = get_embeddings(query)
+    results = client.collections.get("Document").query.near_vector(
+        near_vector=query_embedding.tolist(),
+        limit=3
+    )
+    return results.objects
+
+def generate_response(context, query):
+    """Generate answer using Cohere model based on context and query."""
+    response = cohere_client.generate(
+        model='command',
+        prompt=f"Context: {context}\n\nQuestion: {query}\nAnswer:",
+        max_tokens=100
+    )
+    return response.generations[0].text.strip()
+
+def qa_pipeline(pdf_file, query):
+    """Main pipeline for QA: parse PDF, embed chunks, query Weaviate, and generate answer."""
+    document_text = load_pdf(pdf_file)
+    document_chunks = [document_text[i:i+500] for i in range(0, len(document_text), 500)]
+
+    upload_document_chunks(document_chunks)
+    top_docs = query_answer(query)
+
+    context = ' '.join([doc.properties['content'] for doc in top_docs])
+    answer = generate_response(context, query)
+
+    return context, answer
+
+# Streamlit UI
+st.set_page_config(page_title="Interactive QA Bot", layout="wide")
+
+st.markdown(
+    """
+    <div style="text-align: center; font-size: 28px; font-weight: bold; margin-bottom: 20px; color: #2D3748;">
+        📄 Interactive QA Bot 🔍
+    </div>
+    <p style="text-align: center; font-size: 16px; color: #4A5568;">
+        Upload a PDF document, ask questions, and receive answers based on the document content.
+    </p>
+    <hr style="border: 1px solid #CBD5E0; margin: 20px 0;">
+    """, unsafe_allow_html=True
+)
+
+col1, col2 = st.columns([1, 2])
+
+with col1:
+    pdf_file = st.file_uploader("📁 Upload PDF", type=["pdf"])
+    query = st.text_input("❓ Ask a Question", placeholder="Enter your question here...")
+    submit = st.button("🔍 Submit")
+
+with col2:
+    doc_segments_output = st.empty()
+    answer_output = st.empty()
+
+if submit:
+    if not pdf_file:
+        st.warning("Please upload a PDF file first.")
+    elif not query.strip():
+        st.warning("Please enter a question.")
+    else:
+        with st.spinner("Processing..."):
+            context, answer = qa_pipeline(pdf_file, query)
+            doc_segments_output.text_area("📜 Retrieved Document Segments", context, height=200)
+            answer_output.text_area("💬 Answer", answer, height=80)
+
+# Optional custom CSS for styling
+st.markdown(
+    """
+    <style>
+    body {
+        background-color: #EDF2F7;
+    }
+    .stFileUploader > div > div > input {
+        background-color: #3182CE !important;
+        color: white !important;
+        padding: 8px !important;
+        border-radius: 5px !important;
+    }
+    button {
+        background-color: #3182CE !important;
+        color: white !important;
+        padding: 10px !important;
+        font-size: 16px !important;
+        border-radius: 5px !important;
+        cursor: pointer;
+        border: none !important;
+    }
+    button:hover {
+        background-color: #2B6CB0 !important;
+    }
+    textarea {
+        border: 2px solid #CBD5E0 !important;
+        border-radius: 8px !important;
+        padding: 10px !important;
+        background-color: #FAFAFA !important;
+    }
+    </style>
+    """, unsafe_allow_html=True
+)