File size: 2,928 Bytes
fcbf118 |
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 |
import os
import gradio as gr
import fitz # PyMuPDF
from sentence_transformers import SentenceTransformer
import chromadb
from chromadb.utils import embedding_functions
import openai
# Load GROQ API Key
openai.api_key = os.getenv("GROQ_API_KEY")
openai.api_base = "https://api.groq.com/openai/v1"
# Load embedding model
embedder = SentenceTransformer("all-MiniLM-L6-v2")
# Set up ChromaDB with persistence
persist_path = "./chroma_db"
db = chromadb.Client(chromadb.config.Settings(persist_directory=persist_path))
collection = db.get_or_create_collection("papers")
# Extract text from uploaded PDF
def extract_text_from_pdf(file):
text = ""
doc = fitz.open(stream=file.read(), filetype="pdf")
for page in doc:
text += page.get_text()
return text
# Chunk and store in vector DB
def chunk_and_store(text):
chunks = [text[i:i+500] for i in range(0, len(text), 500)]
embeddings = embedder.encode(chunks).tolist()
for i, chunk in enumerate(chunks):
collection.add(documents=[chunk], ids=[f"id_{len(collection.get()['ids']) + i}"], embeddings=[embeddings[i]])
db.persist()
# Retrieve relevant chunks and send to LLaMA3 via Groq
def retrieve_and_ask(query):
if len(collection.get()["documents"]) == 0:
return "Please upload a paper first."
query_embedding = embedder.encode([query]).tolist()[0]
results = collection.query(query_embeddings=[query_embedding], n_results=3)
context = "\n".join(results["documents"][0])
system_prompt = "You are an academic assistant helping students understand research papers."
user_prompt = f"Based on the following context:\n{context}\n\nAnswer the question:\n{query}"
try:
response = openai.ChatCompletion.create(
model="llama3-70b-8192",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
)
return response['choices'][0]['message']['content']
except Exception as e:
return f"Error: {str(e)}"
# Gradio UI
def handle_upload(file):
if file is None:
return "Upload a valid PDF file."
text = extract_text_from_pdf(file)
chunk_and_store(text)
return "✅ Paper uploaded and processed."
def handle_query(query):
return retrieve_and_ask(query)
with gr.Blocks() as demo:
gr.Markdown("### 📘 RAG Academic Assistant\nUpload a paper and ask questions.")
with gr.Row():
file = gr.File(label="Upload PDF", type="binary")
upload_btn = gr.Button("Process")
upload_output = gr.Textbox()
with gr.Row():
query = gr.Textbox(label="Ask a question")
response = gr.Textbox(label="Answer")
ask_btn = gr.Button("Ask")
upload_btn.click(handle_upload, inputs=[file], outputs=[upload_output])
ask_btn.click(handle_query, inputs=[query], outputs=[response])
demo.launch()
|