Upload 2 files

app.py

@@ -1,142 +1,92 @@
+
 import os
 import fitz  # PyMuPDF
-import numpy as np
 import faiss
+import numpy as np
 import pickle
-import gradio as gr
+import torch
 from sentence_transformers import SentenceTransformer
 from transformers import pipeline
+import gradio as gr
 
-# === Configuration ===
-INDEX_DIR = "saved_index"
-MODEL_NAME = "all-MiniLM-L6-v2"
-CHUNK_SIZE = 500
-
-# === Ensure save directory exists ===
-os.makedirs(INDEX_DIR, exist_ok=True)
-
-# === 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_by_sections(file_path):
-    doc = fitz.open(file_path)
-    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)
-    index = faiss.IndexFlatL2(embeddings.shape[1])
-    index.add(np.array(embeddings))
-    return index, embeddings
-
-def save_index(index, embeddings, chunks):
-    faiss.write_index(index, os.path.join(INDEX_DIR, "faiss.index"))
-    with open(os.path.join(INDEX_DIR, "chunks.pkl"), "wb") as f:
-        pickle.dump(chunks, f)
-    with open(os.path.join(INDEX_DIR, "embeddings.npy"), "wb") as f:
-        np.save(f, embeddings)
-
-def load_index():
-    index_path = os.path.join(INDEX_DIR, "faiss.index")
-    if os.path.exists(index_path):
-        index = faiss.read_index(index_path)
-        with open(os.path.join(INDEX_DIR, "chunks.pkl"), "rb") as f:
-            chunks = pickle.load(f)
-        embeddings = np.load(os.path.join(INDEX_DIR, "embeddings.npy"))
-        return index, embeddings, chunks
-    else:
-        return None, None, []
+# Load or create FAISS index and associated data
+INDEX_FILE = "faiss_index.bin"
+CHUNKS_FILE = "chunks.pkl"
+
+model = SentenceTransformer("all-MiniLM-L6-v2")
+llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", tokenizer="mistralai/Mistral-7B-Instruct-v0.2", device=-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 load_pdf(file):
+    doc = fitz.open(file)
+    text = "\n".join(page.get_text() for page in doc)
+    return text
 
-def ask_question(query, index, embeddings, chunks, top_k=3):
+def split_text(text, chunk_size=500):
+    words = text.split()
+    return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
+
+def create_or_load_index(chunks):
+    if os.path.exists(INDEX_FILE) and os.path.exists(CHUNKS_FILE):
+        with open(CHUNKS_FILE, "rb") as f:
+            chunks = pickle.load(f)
+        index = faiss.read_index(INDEX_FILE)
+    else:
+        embeddings = model.encode(chunks)
+        index = faiss.IndexFlatL2(embeddings.shape[1])
+        index.add(np.array(embeddings))
+        faiss.write_index(index, INDEX_FILE)
+        with open(CHUNKS_FILE, "wb") as f:
+            pickle.dump(chunks, f)
+    return index, chunks
+
+def retrieve_context(query, index, chunks, top_k=3):
     query_emb = model.encode([query])
     distances, indices = index.search(np.array(query_emb), top_k)
-    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": []}
-
-def process_pdfs(files):
-    all_chunks = []
-    for file in files:
-        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"Processed {len(files)} file(s) successfully."
-    else:
-        return "No usable text found in PDFs."
-
-def query_bot(query):
-    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 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"])
+    return "\n\n".join([chunks[i] for i in indices[0]])
+
+def answer_question(query, index, chunks):
+    context = retrieve_context(query, index, chunks)
+    prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
+    response = llm(prompt, max_new_tokens=256, do_sample=False)
+    return response[0]["generated_text"].split("Answer:")[-1].strip()
 
+def generate_quiz(index, chunks):
+    context = retrieve_context("generate quiz questions", index, chunks)
+    prompt = f"Based on the following context, generate 3 quiz questions with multiple choice answers:\n\n{context}\n\nQuestions:"
+    response = llm(prompt, max_new_tokens=512, do_sample=False)
+    return response[0]["generated_text"].split("Questions:")[-1].strip()
+
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🎓 AI Study Assistant (PDF-based with Smart Q&A + Quizzing)")
+    state = {"index": None, "chunks": []}
+
+    gr.Markdown("# 📘 AI Revision Assistant")
+
     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)
+        file_input = gr.File(file_types=[".pdf"], file_count="multiple", label="Upload your revision PDFs")
+        status_output = gr.Textbox(label="Status", interactive=False)
+
+    def process(files):
+        all_chunks = []
+        for file in files:
+            text = load_pdf(file.name)
+            chunks = split_text(text)
+            all_chunks.extend(chunks)
+        index, chunks = create_or_load_index(all_chunks)
+        state["index"] = index
+        state["chunks"] = chunks
+        return f"Processed {len(files)} files. You can now ask questions or generate quizzes."
+
+    file_input.change(fn=process, inputs=file_input, outputs=status_output)
+
+    question_input = gr.Textbox(label="Ask a revision question")
+    answer_output = gr.Textbox(label="Answer", lines=5)
+
+    question_input.submit(fn=lambda q: answer_question(q, state["index"], state["chunks"]) if state["index"] else "Please upload files first.", inputs=question_input, outputs=answer_output)
 
-    gr.Markdown("## ❓ Ask a Question from Your Material")
-    question = gr.Textbox(label="Question")
-    answer = gr.Textbox(label="Answer")
-    question.submit(query_bot, inputs=question, outputs=answer)
+    quiz_btn = gr.Button("Quiz Me")
+    quiz_output = gr.Textbox(label="Generated Quiz Questions", lines=6)
 
-    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)
+    quiz_btn.click(fn=lambda: generate_quiz(state["index"], state["chunks"]) if state["index"] else "Please upload files first.", outputs=quiz_output)
 
-demo.launch()
+demo.launch(debug=True)

requirements.txt

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