Sahil4/docqa-bot

1import gradio as gr2from PyPDF2 import PdfReader3from sentence_transformers import SentenceTransformer4import requests5import numpy as np6import os7import re 8 9HF_API_KEY = os.getenv("HF_API_TOKEN")  10HF_MODEL = "philschmid/bart-large-cnn-samsum"11HF_API_URL = f"https://api-inference.huggingface.co/models/{HF_MODEL}"12HEADERS = {"Authorization": f"Bearer {HF_API_KEY}"}13 14CHUNK_SIZE = 500 15TOP_N_CHUNKS = 1  16MAX_CONTEXT_CHARS = 2000  17 18DOCUMENT_CHUNKS = []19DOCUMENT_EMBEDDINGS = None20conversation_history = []21 22embedder = SentenceTransformer('all-MiniLM-L6-v2')23 24def extract_text_from_pdf(file):25    reader = PdfReader(file.name)26    text = ""27    for page in reader.pages:28        page_text = page.extract_text()29        if page_text:30            text += page_text + " "31    if not text.strip():32        return None33    return text34 35def chunk_text(text, chunk_size=CHUNK_SIZE):36    chunks = []37    for i in range(0, len(text), chunk_size):38        chunk = text[i:i+chunk_size].strip()39        if chunk:40            chunks.append(chunk)41    print(f"[INFO] Total chunks created: {len(chunks)}")42    return chunks43 44def process_document(file):45    global DOCUMENT_CHUNKS, DOCUMENT_EMBEDDINGS, conversation_history46    text = extract_text_from_pdf(file)47    if not text:48        return "Uploaded PDF is empty or unreadable."49    50    DOCUMENT_CHUNKS = chunk_text(text)51    DOCUMENT_EMBEDDINGS = embedder.encode(DOCUMENT_CHUNKS)52    conversation_history = []53    54    return f"Document processed successfully! Total chunks: {len(DOCUMENT_CHUNKS)}"55 56def retrieve_chunks(question, top_n=TOP_N_CHUNKS):57    if not DOCUMENT_CHUNKS or DOCUMENT_EMBEDDINGS is None:58        return []59    question_emb = embedder.encode([question])60    sims = [np.dot(question_emb[0], emb) / (np.linalg.norm(question_emb[0]) * np.linalg.norm(emb))61            for emb in DOCUMENT_EMBEDDINGS]62    top_indices = np.argsort(sims)[::-1][:min(top_n, len(DOCUMENT_CHUNKS))]63    top_chunks = [DOCUMENT_CHUNKS[i] for i in top_indices]64    print(f"[INFO] Top chunk indices: {top_indices}")65    return top_chunks66 67 68def ask_question(user_question):69    global conversation_history70    71    if not DOCUMENT_CHUNKS:72        return [["System", "Please upload a document first."]]73    74    greeting_keywords = ["hi", "hello", "hey", "good morning", "good afternoon"]75    thanks_keywords = ["thanks", "thank you", "thx", "thank u"]76    77    user_lower = user_question.strip().lower()78    if any(re.search(rf"\b{word}\b", user_lower) for word in greeting_keywords):79        return [["System", "Hi! How can I help you?"]]80 81    if any(re.search(rf"\b{word}\b", user_lower) for word in thanks_keywords):82        return [["System", "You're welcome!"]]83    84    q_emb = embedder.encode([user_question])[0]85    for qa in conversation_history[-10:]:86        prev_q, prev_a, prev_emb = qa87        score = np.dot(q_emb, prev_emb) / (np.linalg.norm(q_emb) * np.linalg.norm(prev_emb))88        if score > 0.85:89            return [[prev_q, prev_a]]  # repeated/related question90    91    relevant_chunks = retrieve_chunks(user_question)92    if not relevant_chunks:93        answer = "No relevant information found in the document."94    else:95        context = "\n".join(relevant_chunks)96        if len(context) > MAX_CONTEXT_CHARS:97            context = context[:MAX_CONTEXT_CHARS]98        99        history_text = ""100        for q, a, _ in conversation_history[-3:]:101            history_text += f"Previous Q: {q}\nPrevious A: {a}\n"102        103        prompt = f"{history_text}Answer based only on the following context:\n{context}\n\nQuestion: {user_question}\nAnswer:"104        105        try:106            response = requests.post(HF_API_URL, headers=HEADERS, json={"inputs": prompt}, timeout=30)107            res_json = response.json()108            if isinstance(res_json, list) and "summary_text" in res_json[0]:109                answer = res_json[0]["summary_text"]110            elif isinstance(res_json, list) and "generated_text" in res_json[0]:111                answer = res_json[0]["generated_text"]112            else:113                answer = f"Unexpected response: {res_json}"114        except requests.exceptions.Timeout:115            answer = "I could not understand your query. Please ask again."116        except Exception as e:117            answer = f"Error: {str(e)}"118    119    conversation_history.append((user_question, answer, q_emb))120    if len(conversation_history) > 10:121        conversation_history.pop(0)122    123    # Format for Gradio Chatbot124    chat_display = [[q, a] for q, a, _ in conversation_history]125    return chat_display126 127 128with gr.Blocks() as demo:129    gr.Markdown("## Hi 👋! Upload a PDF and ask questions about it.")130    131    with gr.Row():132        pdf_input = gr.File(label="Upload PDF", file_types=['.pdf'])133        upload_btn = gr.Button("Process Document")134    135    status = gr.Textbox(label="Status", interactive=False)136    137    question = gr.Textbox(label="Ask a question")138    ask_btn = gr.Button("Ask")139    chatbot = gr.Chatbot(label="Chat History")140    141    upload_btn.click(fn=process_document, inputs=pdf_input, outputs=status)142    ask_btn.click(fn=ask_question, inputs=question, outputs=chatbot)143 144demo.launch()