hmdsolution/DocumentAIAgent

1import os2import gradio as gr3import faiss4import numpy as np5from sentence_transformers import SentenceTransformer6from groq import Groq7from PyPDF2 import PdfReader8 9# ==============================10# CONFIG11# ==============================12 13GROQ_API_KEY = os.getenv("GROQ_API_KEY")14client = Groq(api_key=GROQ_API_KEY)15 16embedding_model = SentenceTransformer("all-MiniLM-L6-v2")17 18index = None19document_chunks = []20 21SIMILARITY_THRESHOLD = 0.65  # strict control22 23 24# ==============================25# DOCUMENT PROCESSING26# ==============================27 28def process_document(file):29    global index, document_chunks30    31    text = ""32    33    if file.name.endswith(".pdf"):34        reader = PdfReader(file.name)35        for page in reader.pages:36            text += page.extract_text()37    else:38        text = file.read().decode("utf-8")39 40    # Chunking41    document_chunks = [text[i:i+500] for i in range(0, len(text), 500)]42 43    embeddings = embedding_model.encode(document_chunks)44    embeddings = np.array(embeddings).astype("float32")45 46    index = faiss.IndexFlatL2(embeddings.shape[1])47    index.add(embeddings)48 49    return "Document processed successfully."50 51 52# ==============================53# CHAT FUNCTION54# ==============================55 56def chat(question):57    global index, document_chunks58 59    if index is None:60        return "Please upload a document first."61 62    question_embedding = embedding_model.encode([question])63    question_embedding = np.array(question_embedding).astype("float32")64 65    D, I = index.search(question_embedding, k=1)66 67    similarity_score = 1 / (1 + D[0][0])68 69    if similarity_score < SIMILARITY_THRESHOLD:70        return "Response not found"71 72    relevant_text = document_chunks[I[0][0]]73 74    prompt = f"""75You are a strict document assistant.76 77Answer ONLY using the provided document context.78If answer is not clearly present, reply exactly:79Response not found80 81Document Context:82{relevant_text}83 84Question:85{question}86 87Answer:88"""89 90    completion = client.chat.completions.create(91        model="deepseek-r1-distill-llama-70b",92        messages=[{"role": "user", "content": prompt}],93        temperature=094    )95 96    return completion.choices[0].message.content.strip()97 98 99# ==============================100# UI101# ==============================102 103with gr.Blocks() as demo:104    gr.Markdown("# 📄 Restricted Document AI Chatbot")105 106    file_input = gr.File(label="Upload Document (PDF or TXT)")107    upload_btn = gr.Button("Process Document")108    status = gr.Textbox()109 110    question = gr.Textbox(label="Ask Question")111    answer = gr.Textbox(label="Answer")112 113    upload_btn.click(process_document, inputs=file_input, outputs=status)114    question.submit(chat, inputs=question, outputs=answer)115 116demo.launch()