agnixcode commited on
Commit
fcbf118
·
verified ·
1 Parent(s): ed124b7

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +89 -0
app.py ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import gradio as gr
3
+ import fitz # PyMuPDF
4
+ from sentence_transformers import SentenceTransformer
5
+ import chromadb
6
+ from chromadb.utils import embedding_functions
7
+ import openai
8
+
9
+ # Load GROQ API Key
10
+ openai.api_key = os.getenv("GROQ_API_KEY")
11
+ openai.api_base = "https://api.groq.com/openai/v1"
12
+
13
+ # Load embedding model
14
+ embedder = SentenceTransformer("all-MiniLM-L6-v2")
15
+
16
+ # Set up ChromaDB with persistence
17
+ persist_path = "./chroma_db"
18
+ db = chromadb.Client(chromadb.config.Settings(persist_directory=persist_path))
19
+ collection = db.get_or_create_collection("papers")
20
+
21
+ # Extract text from uploaded PDF
22
+ def extract_text_from_pdf(file):
23
+ text = ""
24
+ doc = fitz.open(stream=file.read(), filetype="pdf")
25
+ for page in doc:
26
+ text += page.get_text()
27
+ return text
28
+
29
+ # Chunk and store in vector DB
30
+ def chunk_and_store(text):
31
+ chunks = [text[i:i+500] for i in range(0, len(text), 500)]
32
+ embeddings = embedder.encode(chunks).tolist()
33
+
34
+ for i, chunk in enumerate(chunks):
35
+ collection.add(documents=[chunk], ids=[f"id_{len(collection.get()['ids']) + i}"], embeddings=[embeddings[i]])
36
+ db.persist()
37
+
38
+ # Retrieve relevant chunks and send to LLaMA3 via Groq
39
+ def retrieve_and_ask(query):
40
+ if len(collection.get()["documents"]) == 0:
41
+ return "Please upload a paper first."
42
+
43
+ query_embedding = embedder.encode([query]).tolist()[0]
44
+ results = collection.query(query_embeddings=[query_embedding], n_results=3)
45
+ context = "\n".join(results["documents"][0])
46
+
47
+ system_prompt = "You are an academic assistant helping students understand research papers."
48
+ user_prompt = f"Based on the following context:\n{context}\n\nAnswer the question:\n{query}"
49
+
50
+ try:
51
+ response = openai.ChatCompletion.create(
52
+ model="llama3-70b-8192",
53
+ messages=[
54
+ {"role": "system", "content": system_prompt},
55
+ {"role": "user", "content": user_prompt}
56
+ ]
57
+ )
58
+ return response['choices'][0]['message']['content']
59
+ except Exception as e:
60
+ return f"Error: {str(e)}"
61
+
62
+ # Gradio UI
63
+ def handle_upload(file):
64
+ if file is None:
65
+ return "Upload a valid PDF file."
66
+ text = extract_text_from_pdf(file)
67
+ chunk_and_store(text)
68
+ return "✅ Paper uploaded and processed."
69
+
70
+ def handle_query(query):
71
+ return retrieve_and_ask(query)
72
+
73
+ with gr.Blocks() as demo:
74
+ gr.Markdown("### 📘 RAG Academic Assistant\nUpload a paper and ask questions.")
75
+
76
+ with gr.Row():
77
+ file = gr.File(label="Upload PDF", type="binary")
78
+ upload_btn = gr.Button("Process")
79
+ upload_output = gr.Textbox()
80
+
81
+ with gr.Row():
82
+ query = gr.Textbox(label="Ask a question")
83
+ response = gr.Textbox(label="Answer")
84
+ ask_btn = gr.Button("Ask")
85
+
86
+ upload_btn.click(handle_upload, inputs=[file], outputs=[upload_output])
87
+ ask_btn.click(handle_query, inputs=[query], outputs=[response])
88
+
89
+ demo.launch()