File size: 3,248 Bytes
d8a0579 |
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 102 103 |
import os
import gradio as gr
import faiss
import numpy as np
import PyPDF2
from sentence_transformers import SentenceTransformer
import requests
# 🔐 Use environment variable for Groq API key
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
GROQ_MODEL = "llama3-8b-8192"
# 🧠 Embedding Model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
# 🗃️ Vector DB
dimension = 384
index = faiss.IndexFlatL2(dimension)
text_chunks = []
def extract_text_from_pdf(pdf_file):
reader = PyPDF2.PdfReader(pdf_file)
text = ""
for page in reader.pages:
text += page.extract_text() or ""
return text
def chunk_text(text, chunk_size=500):
sentences = text.split('. ')
chunks, chunk = [], ""
for sentence in sentences:
if len(chunk) + len(sentence) < chunk_size:
chunk += sentence + ". "
else:
chunks.append(chunk.strip())
chunk = sentence + ". "
if chunk:
chunks.append(chunk.strip())
return chunks
def embed_and_store(chunks):
global text_chunks
text_chunks = chunks
embeddings = embedding_model.encode(chunks)
index.add(np.array(embeddings))
def retrieve_context(query, top_k=3):
query_vector = embedding_model.encode([query])
distances, indices = index.search(np.array(query_vector), top_k)
return "\n".join([text_chunks[i] for i in indices[0]])
def format_prompt(context, question):
system_msg = "You are a helpful research assistant who answers questions using only the uploaded document."
user_msg = f"Document Context:\n{context}\n\nQuestion: {question}\nAnswer:"
return [{"role": "system", "content": system_msg},
{"role": "user", "content": user_msg}]
def call_groq_api(messages):
url = "https://api.groq.com/openai/v1/chat/completions"
headers = {
"Authorization": f"Bearer {GROQ_API_KEY}",
"Content-Type": "application/json"
}
data = {
"model": GROQ_MODEL,
"messages": messages,
"temperature": 0.3
}
response = requests.post(url, headers=headers, json=data)
return response.json()['choices'][0]['message']['content']
def upload_file(pdf):
text = extract_text_from_pdf(pdf)
chunks = chunk_text(text)
embed_and_store(chunks)
return "✅ Document processed. You may now ask questions."
def answer_question(question):
if not text_chunks:
return "❌ Please upload and process a document first."
context = retrieve_context(question)
messages = format_prompt(context, question)
return call_groq_api(messages)
with gr.Blocks() as rag_ui:
gr.Markdown("## 📄 RAG Assistant with LLaMA3 (Groq)")
with gr.Row():
pdf_input = gr.File(label="Upload PDF")
upload_button = gr.Button("Process Document")
status_output = gr.Textbox(label="Status")
upload_button.click(upload_file, inputs=pdf_input, outputs=status_output)
gr.Markdown("### ❓ Ask a Question from the Uploaded PDF")
question_input = gr.Textbox(label="Your Question")
answer_output = gr.Textbox(label="Answer", lines=5)
ask_button = gr.Button("Get Answer")
ask_button.click(answer_question, inputs=question_input, outputs=answer_output)
rag_ui.launch()
|