Spaces:
Build error
Build error
File size: 3,637 Bytes
b7b4342 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import os
import fitz # PyMuPDF
import gradio as gr
import requests
from dotenv import load_dotenv
from sentence_transformers import SentenceTransformer, util
# Load environment variables
load_dotenv()
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
MODEL_NAME = "llama3-8b-8192"
# Load sentence-transformer model
embed_model = SentenceTransformer('all-MiniLM-L6-v2')
SYSTEM_PROMPT = """
You are a helpful AI assistant that answers questions using provided PDF documents.
Always base your answers on the context and include page numbers when possible.
"""
# Extract text from PDF with page numbers
def extract_text_with_pages(file):
doc = fitz.open(file.name)
chunks = []
for page_num in range(len(doc)):
page = doc.load_page(page_num)
text = page.get_text()
if text.strip():
chunks.append({"text": text.strip(), "page": page_num + 1})
return chunks
# Chunking (simple, per page)
def prepare_embeddings(chunks):
texts = [chunk["text"] for chunk in chunks]
embeddings = embed_model.encode(texts, convert_to_tensor=True)
return embeddings
# Retrieve top chunks
def retrieve_top_chunks(query, chunks, embeddings, top_k=3):
query_embedding = embed_model.encode(query, convert_to_tensor=True)
hits = util.semantic_search(query_embedding, embeddings, top_k=top_k)[0]
results = []
for hit in hits:
matched_chunk = chunks[hit["corpus_id"]]
results.append(f"[Page {matched_chunk['page']}]: {matched_chunk['text']}")
return results
# Query GROQ LLM
def query_groq(question, context, chat_history):
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
for user, bot in chat_history:
messages.append({"role": "user", "content": user})
messages.append({"role": "assistant", "content": bot})
messages.append({"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"})
headers = {
"Authorization": f"Bearer {GROQ_API_KEY}",
"Content-Type": "application/json"
}
response = requests.post(GROQ_API_URL, headers=headers, json={
"model": MODEL_NAME,
"messages": messages,
"temperature": 0.7
})
if response.status_code == 200:
return response.json()["choices"][0]["message"]["content"]
else:
return f"Error {response.status_code}: {response.text}"
# Full RAG pipeline
def rag_chatbot(files, question, history):
all_chunks = []
for file in files:
all_chunks.extend(extract_text_with_pages(file))
embeddings = prepare_embeddings(all_chunks)
top_chunks = retrieve_top_chunks(question, all_chunks, embeddings)
context = "\n\n".join(top_chunks)
response = query_groq(question, context, history)
history.append((question, response))
return "", history
# Gradio Interface
with gr.Blocks() as demo:
gr.Markdown("## 📄 RAG PDF Chatbot (Groq + SentenceTransformers)")
gr.Markdown("Upload PDFs and ask questions. The assistant will use document context and show page numbers in answers.")
pdf_files = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload your PDF files")
question = gr.Textbox(label="Ask a question about the documents")
chatbot = gr.Chatbot()
clear_btn = gr.Button("Clear Chat")
state = gr.State([])
question.submit(rag_chatbot, [pdf_files, question, state], [question, chatbot])
clear_btn.click(lambda: ([], []), None, [chatbot, state])
demo.launch()
|