app.py

import os
import gradio as gr
import PyPDF2
import torch
import weaviate
from transformers import AutoTokenizer, AutoModel
from weaviate.classes.init import Auth
import cohere

# Load credentials from environment variables
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)
    return ''.join([page.extract_text() for page in reader.pages if page.extract_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

# Gradio UI
with gr.Blocks(theme="compact") as demo:
    gr.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;">
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            pdf_input = gr.File(label="📁 Upload PDF", file_types=[".pdf"])
            query_input = gr.Textbox(label="❓ Ask a Question", placeholder="Enter your question here...")
            submit_button = gr.Button("🔍 Submit")

        with gr.Column(scale=2):
            doc_segments_output = gr.Textbox(label="📜 Retrieved Document Segments", lines=10)
            answer_output = gr.Textbox(label="💬 Answer", lines=3)

    submit_button.click(
        qa_pipeline,
        inputs=[pdf_input, query_input],
        outputs=[doc_segments_output, answer_output]
    )

    gr.Markdown(
        """
        <style>
            body {
                background-color: #EDF2F7;
            }
            input[type="file"] {
                background-color: #3182CE;
                color: white;
                padding: 8px;
                border-radius: 5px;
            }
            button {
                background-color: #3182CE;
                color: white;
                padding: 10px;
                font-size: 16px;
                border-radius: 5px;
                cursor: pointer;
            }
            button:hover {
                background-color: #2B6CB0;
            }
            textarea {
                border: 2px solid #CBD5E0;
                border-radius: 8px;
                padding: 10px;
                background-color: #FAFAFA;
            }
        </style>
        """
    )

demo.launch()