app.py

# app.py

import gradio as gr
import os
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer
from groq import Groq
from docx import Document

# -------------------- API Keys --------------------
groq_api_key = os.getenv("GROQ_API_KEY")
client = Groq(api_key=groq_api_key)

# -------------------- Models --------------------
embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')

# -------------------- FAISS Setup --------------------
faiss_index = faiss.IndexFlatL2(384)
document_chunks = []

# -------------------- Utils --------------------
def read_docx(file_obj):
    doc = Document(file_obj)
    return "\n".join([para.text for para in doc.paragraphs])

def chunk_text(text, chunk_size=500, overlap=50):
    tokens = tokenizer.tokenize(text)
    chunks = []
    for i in range(0, len(tokens), chunk_size - overlap):
        chunk = tokens[i:i+chunk_size]
        chunk_text = tokenizer.convert_tokens_to_string(chunk)
        chunks.append(chunk_text)
    return chunks

def embed_chunks(chunks):
    return embedding_model.encode(chunks, convert_to_numpy=True)

def process_document(doc):
    if not doc.name.endswith(".docx"):
        return "Only .docx files are supported."

    text = read_docx(doc)
    if not text.strip():
        return "The document is empty."

    chunks = chunk_text(text)
    if not chunks:
        return "No valid chunks generated from the document."

    embeddings = embed_chunks(chunks)
    faiss_index.add(embeddings)
    document_chunks.extend(chunks)

    return f"Document processed and {len(chunks)} chunks indexed."

def retrieve_relevant_chunks(query, top_k=3):
    if faiss_index.ntotal == 0:
        return ["No document indexed. Please upload and process a document first."]

    query_vec = embedding_model.encode([query], convert_to_numpy=True)
    D, I = faiss_index.search(query_vec, top_k)
    return [document_chunks[i] for i in I[0] if i < len(document_chunks)]

def generate_answer_with_groq(prompt):
    try:
        chat_completion = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2,
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"Error from Groq API: {e}"

def answer_query(query):
    if faiss_index.ntotal == 0 or len(document_chunks) == 0:
        return "No document indexed. Please upload and process a document first."

    context = retrieve_relevant_chunks(query)
    combined_context = "\n".join(context)
    prompt = f"Context:\n{combined_context}\n\nQuestion: {query}"
    return generate_answer_with_groq(prompt)

# -------------------- Gradio UI --------------------
with gr.Blocks() as rag_app:
    gr.Markdown("## 🧠 RAG System with .docx Upload | Hugging Face + Groq")

    with gr.Row():
        upload = gr.File(label="Upload a .docx file", file_types=[".docx"])
        upload_button = gr.Button("Process Document")

    query_input = gr.Textbox(label="Ask a question")
    query_button = gr.Button("Submit")
    output = gr.Textbox(label="Answer")

    upload_button.click(fn=process_document, inputs=upload, outputs=output)
    query_button.click(fn=answer_query, inputs=query_input, outputs=output)

rag_app.launch()