Upload 2 files

app.py

@@ -2,9 +2,10 @@ import os
 import fitz  # PyMuPDF
 import numpy as np
 import faiss
+import pickle
 import gradio as gr
 from sentence_transformers import SentenceTransformer
-import pickle
+from transformers import pipeline
 
 # === Configuration ===
 INDEX_DIR = "saved_index"
@@ -14,17 +15,31 @@ CHUNK_SIZE = 500
 # === Ensure save directory exists ===
 os.makedirs(INDEX_DIR, exist_ok=True)
 
-# === Load model globally ===
+# === Load sentence transformer and LLM pipeline ===
 model = SentenceTransformer(MODEL_NAME)
+llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", tokenizer="mistralai/Mistral-7B-Instruct-v0.2", device=0 if torch.cuda.is_available() else -1)
 
 # === Helper functions ===
-def load_pdf(file_path):
+def load_pdf_by_sections(file_path):
     doc = fitz.open(file_path)
-    return "\n".join(page.get_text() for page in doc)
-
-def split_into_chunks(text, chunk_size=CHUNK_SIZE):
-    words = text.split()
-    return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
+    chunks = []
+    for page in doc:
+        blocks = page.get_text("blocks")
+        blocks.sort()  # sort top-to-bottom
+        text = ""
+        for b in blocks:
+            content = b[4].strip()
+            if len(content.split()) < 5:
+                continue
+            if content.isupper() or content.istitle():  # crude heading detection
+                if text:
+                    chunks.append(text.strip())
+                text = content
+            else:
+                text += " " + content
+        if text:
+            chunks.append(text.strip())
+    return chunks
 
 def create_index(chunks):
     embeddings = model.encode(chunks)
@@ -50,15 +65,28 @@ def load_index():
     else:
         return None, None, []
 
-def ask_question(query, index, embeddings, chunks, top_k=3, context_window=1):
+def generate_answer(context, question):
+    prompt = f"Answer this based only on the context below. Be precise and relevant.\n\nContext:\n{context}\n\nQuestion: {question}\nAnswer:"
+    response = llm(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
+    return response.split("Answer:")[-1].strip()
+
+def ask_question(query, index, embeddings, chunks, top_k=3):
     query_emb = model.encode([query])
     distances, indices = index.search(np.array(query_emb), top_k)
-    context_chunks = set()
-    for i in indices[0]:
-        for j in range(max(0, i - context_window), min(len(chunks), i + context_window + 1)):
-            context_chunks.add(j)
-    sorted_context = sorted(list(context_chunks))
-    return "\n\n".join([chunks[i] for i in sorted_context])
+    relevant_context = "\n\n".join([chunks[i] for i in indices[0]])
+    return generate_answer(relevant_context, query)
+
+def quiz_me(index, embeddings, chunks):
+    np.random.seed(42)
+    sample_indices = np.random.choice(len(chunks), 2, replace=False)
+    questions = []
+    for i in sample_indices:
+        context = chunks[i]
+        prompt = f"Based only on the text below, generate a single clear quiz question:\n\n{context}\n\nQuestion:"
+        result = llm(prompt, max_new_tokens=100, do_sample=True, temperature=0.7)[0]["generated_text"]
+        extracted = result.split("Question:")[-1].strip()
+        questions.append(extracted)
+    return "\n\n".join(questions)
 
 # === Gradio Interface ===
 state = {"index": None, "embeddings": None, "chunks": []}
@@ -66,16 +94,15 @@ state = {"index": None, "embeddings": None, "chunks": []}
 def process_pdfs(files):
     all_chunks = []
     for file in files:
-        text = load_pdf(file.name)
-        chunks = split_into_chunks(text)
+        chunks = load_pdf_by_sections(file.name)
         all_chunks.extend(chunks)
     if all_chunks:
         index, embeddings = create_index(all_chunks)
         save_index(index, embeddings, all_chunks)
         state.update({"index": index, "embeddings": embeddings, "chunks": all_chunks})
-        return f"Uploaded and indexed {len(files)} file(s) successfully."
+        return f"Processed {len(files)} file(s) successfully."
     else:
-        return "No text found in PDFs."
+        return "No usable text found in PDFs."
 
 def query_bot(query):
     if not state["index"]:
@@ -86,14 +113,30 @@ def query_bot(query):
             return "Please upload and process PDFs first."
     return ask_question(query, state["index"], state["embeddings"], state["chunks"])
 
+def quiz_bot(dummy=""):
+    if not state["index"]:
+        index, embeddings, chunks = load_index()
+        if index:
+            state.update({"index": index, "embeddings": embeddings, "chunks": chunks})
+        else:
+            return "Please upload and process PDFs first."
+    return quiz_me(state["index"], state["embeddings"], state["chunks"])
+
 with gr.Blocks() as demo:
-    gr.Markdown("# 📘 AI Revision Assistant – Uses Only Your PDFs")
+    gr.Markdown("# 🎓 AI Study Assistant (PDF-based with Smart Q&A + Quizzing)")
     with gr.Row():
         file_input = gr.File(label="Upload PDFs", file_count="multiple")
         status = gr.Textbox(label="Status Message")
     file_input.change(process_pdfs, inputs=file_input, outputs=status)
-    query = gr.Textbox(label="Ask a Question")
+
+    gr.Markdown("## ❓ Ask a Question from Your Material")
+    question = gr.Textbox(label="Question")
     answer = gr.Textbox(label="Answer")
-    query.submit(query_bot, inputs=query, outputs=answer)
+    question.submit(query_bot, inputs=question, outputs=answer)
+
+    gr.Markdown("## 🧠 Quiz Me from Uploaded PDFs")
+    quiz_btn = gr.Button("Generate 2 Quiz Questions")
+    quiz_output = gr.Textbox(label="Quiz Questions")
+    quiz_btn.click(quiz_bot, inputs=[], outputs=quiz_output)
 
 demo.launch()

requirements.txt

@@ -1,4 +1,8 @@
 pymupdf
 sentence-transformers
 faiss-cpu
-gradio
+gradio
+transformers
+torch
+accelerate
+safetensors