Create app.py

app.py

@@ -0,0 +1,528 @@
+import os
+import re
+import json
+import pickle
+from urllib.parse import quote
+
+import numpy as np
+import gradio as gr
+from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer
+from openai import OpenAI
+
+# =====================================================
+# PATHS
+# =====================================================
+BUILD_DIR = "brainchat_build"
+CHUNKS_PATH = os.path.join(BUILD_DIR, "chunks.pkl")
+TOKENS_PATH = os.path.join(BUILD_DIR, "tokenized_chunks.pkl")
+EMBED_PATH = os.path.join(BUILD_DIR, "embeddings.npy")
+CONFIG_PATH = os.path.join(BUILD_DIR, "config.json")
+LOGO_FILE = "Brain chat-09.png"
+
+# =====================================================
+# GLOBALS
+# =====================================================
+EMBED_MODEL = None
+BM25 = None
+CHUNKS = None
+EMBEDDINGS = None
+CLIENT = None
+
+
+# =====================================================
+# LOADERS
+# =====================================================
+def tokenize(text: str):
+    return re.findall(r"\w+", text.lower(), flags=re.UNICODE)
+
+
+def ensure_loaded():
+    global EMBED_MODEL, BM25, CHUNKS, EMBEDDINGS, CLIENT
+
+    if CHUNKS is None:
+        for path in [CHUNKS_PATH, TOKENS_PATH, EMBED_PATH, CONFIG_PATH]:
+            if not os.path.exists(path):
+                raise FileNotFoundError(f"Missing file: {path}")
+
+        with open(CHUNKS_PATH, "rb") as f:
+            CHUNKS = pickle.load(f)
+
+        with open(TOKENS_PATH, "rb") as f:
+            tokenized_chunks = pickle.load(f)
+
+        EMBEDDINGS = np.load(EMBED_PATH)
+
+        with open(CONFIG_PATH, "r", encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        BM25 = BM25Okapi(tokenized_chunks)
+        EMBED_MODEL = SentenceTransformer(cfg["embedding_model"])
+
+    if CLIENT is None:
+        api_key = os.getenv("OPENAI_API_KEY")
+        if not api_key:
+            raise ValueError("OPENAI_API_KEY is missing in Hugging Face Space Secrets.")
+        CLIENT = OpenAI(api_key=api_key)
+
+
+# =====================================================
+# RETRIEVAL
+# =====================================================
+def search_hybrid(query: str, shortlist_k: int = 20, final_k: int = 3):
+    ensure_loaded()
+
+    query_tokens = tokenize(query)
+    bm25_scores = BM25.get_scores(query_tokens)
+
+    shortlist_idx = np.argsort(bm25_scores)[::-1][:shortlist_k]
+    shortlist_embeddings = EMBEDDINGS[shortlist_idx]
+
+    qvec = EMBED_MODEL.encode([query], normalize_embeddings=True).astype("float32")[0]
+    dense_scores = shortlist_embeddings @ qvec
+
+    rerank_order = np.argsort(dense_scores)[::-1][:final_k]
+    final_idx = shortlist_idx[rerank_order]
+
+    return [CHUNKS[int(i)] for i in final_idx]
+
+
+def build_context(records):
+    blocks = []
+    for i, r in enumerate(records, start=1):
+        blocks.append(
+            f"""[Source {i}]
+Book: {r['book']}
+Section: {r['section_title']}
+Pages: {r['page_start']}-{r['page_end']}
+Text:
+{r['text']}"""
+        )
+    return "\n\n".join(blocks)
+
+
+def make_sources(records):
+    seen = set()
+    lines = []
+    for r in records:
+        key = (r["book"], r["section_title"], r["page_start"], r["page_end"])
+        if key in seen:
+            continue
+        seen.add(key)
+        lines.append(
+            f"• {r['book']} | {r['section_title']} | pp. {r['page_start']}-{r['page_end']}"
+        )
+    return "\n".join(lines)
+
+
+# =====================================================
+# PROMPTS
+# =====================================================
+def language_instruction(language_mode: str) -> str:
+    if language_mode == "English":
+        return "Answer only in English."
+    if language_mode == "Spanish":
+        return "Answer only in Spanish."
+    if language_mode == "Bilingual":
+        return "Answer first in English, then provide a Spanish version under the heading 'Español:'."
+    return (
+        "If the user's message is in Spanish, answer in Spanish. "
+        "If the user's message is in English, answer in English."
+    )
+
+
+def choose_quiz_count(user_text: str, selector: str) -> int:
+    if selector in {"3", "5", "7"}:
+        return int(selector)
+
+    t = user_text.lower()
+    if any(k in t for k in ["mock test", "final exam", "exam practice", "full test"]):
+        return 7
+    if any(k in t for k in ["detailed", "revision", "comprehensive", "study"]):
+        return 5
+    return 3
+
+
+def build_tutor_prompt(mode: str, language_mode: str, question: str, context: str) -> str:
+    mode_map = {
+        "Explain": "Explain clearly like a friendly tutor using simple language.",
+        "Detailed": "Give a fuller and more detailed explanation.",
+        "Short Notes": "Answer in concise revision-note format using bullets.",
+        "Flashcards": "Create 6 flashcards in Q/A format.",
+        "Case-Based": "Create a short clinical scenario and explain it clearly."
+    }
+
+    return f"""
+You are BrainChat, an interactive neurology and neuroanatomy tutor.
+
+Rules:
+- Use only the provided context from the books.
+- If the answer is not supported by the context, say exactly:
+  Not found in the course material.
+- Be accurate and student-friendly.
+- {language_instruction(language_mode)}
+
+Teaching style:
+{mode_map[mode]}
+
+Context:
+{context}
+
+Question:
+{question}
+""".strip()
+
+
+def build_quiz_generation_prompt(language_mode: str, topic: str, context: str, n_questions: int) -> str:
+    return f"""
+You are BrainChat, an interactive tutor.
+
+Rules:
+- Use only the provided context.
+- Create exactly {n_questions} quiz questions.
+- Questions should be short and clear.
+- Also create a short answer key.
+- Return valid JSON only.
+- {language_instruction(language_mode)}
+
+Required JSON format:
+{{
+  "title": "short quiz title",
+  "questions": [
+    {{"q": "question 1", "answer_key": "expected short answer"}},
+    {{"q": "question 2", "answer_key": "expected short answer"}}
+  ]
+}}
+
+Context:
+{context}
+
+Topic:
+{topic}
+""".strip()
+
+
+def build_quiz_evaluation_prompt(language_mode: str, quiz_data: dict, user_answers: str) -> str:
+    quiz_json = json.dumps(quiz_data, ensure_ascii=False)
+    return f"""
+You are BrainChat, an interactive tutor.
+
+Evaluate the student's answers fairly using the quiz answer key.
+Give:
+- total score
+- per-question feedback
+- one short improvement suggestion
+
+Rules:
+- Accept semantically correct answers even if wording differs.
+- Return valid JSON only.
+- {language_instruction(language_mode)}
+
+Required JSON format:
+{{
+  "score_obtained": 0,
+  "score_total": 0,
+  "summary": "short overall feedback",
+  "results": [
+    {{
+      "question": "question text",
+      "student_answer": "student answer",
+      "result": "Correct / Partially Correct / Incorrect",
+      "feedback": "short explanation"
+    }}
+  ]
+}}
+
+Quiz data:
+{quiz_json}
+
+Student answers:
+{user_answers}
+""".strip()
+
+
+def chat_text(prompt: str) -> str:
+    resp = CLIENT.chat.completions.create(
+        model="gpt-4o-mini",
+        temperature=0.2,
+        messages=[
+            {"role": "system", "content": "You are a helpful educational assistant."},
+            {"role": "user", "content": prompt},
+        ],
+    )
+    return resp.choices[0].message.content.strip()
+
+
+def chat_json(prompt: str) -> dict:
+    resp = CLIENT.chat.completions.create(
+        model="gpt-4o-mini",
+        temperature=0.2,
+        response_format={"type": "json_object"},
+        messages=[
+            {"role": "system", "content": "Return only valid JSON."},
+            {"role": "user", "content": prompt},
+        ],
+    )
+    return json.loads(resp.choices[0].message.content)
+
+
+# =====================================================
+# UI HELPERS
+# =====================================================
+def detect_logo_url():
+    if os.path.exists(LOGO_FILE):
+        return f"/gradio_api/file={quote(LOGO_FILE)}"
+    return None
+
+
+def render_header():
+    logo_url = detect_logo_url()
+    if logo_url:
+        logo_html = f'<img src="{logo_url}" alt="BrainChat Logo" style="width:120px;height:120px;object-fit:contain;display:block;margin:0 auto;">'
+    else:
+        logo_html = '<div style="width:120px;height:120px;border-radius:50%;background:#efe85a;display:flex;align-items:center;justify-content:center;font-weight:700;text-align:center;margin:0 auto;">BRAIN<br>CHAT</div>'
+
+    return f"""
+    <div class="hero-card">
+        <div class="hero-inner">
+            <div class="hero-logo">{logo_html}</div>
+            <div class="hero-title">BrainChat</div>
+            <div class="hero-subtitle">
+                Interactive neurology and neuroanatomy tutor based on your uploaded books
+            </div>
+        </div>
+    </div>
+    """
+
+
+# =====================================================
+# MAIN CHAT FUNCTION
+# =====================================================
+def answer_question(message, history, mode, language_mode, quiz_count_mode, show_sources, quiz_state):
+    if history is None:
+        history = []
+    if quiz_state is None:
+        quiz_state = {
+            "active": False,
+            "topic": None,
+            "quiz_data": None,
+            "language_mode": "Auto"
+        }
+
+    if not message or not message.strip():
+        return history, quiz_state, ""
+
+    try:
+        ensure_loaded()
+        user_text = message.strip()
+
+        # Add user message first
+        history.append({"role": "user", "content": user_text})
+
+        # -----------------------------------------
+        # QUIZ EVALUATION STEP
+        # -----------------------------------------
+        if quiz_state.get("active", False):
+            evaluation_prompt = build_quiz_evaluation_prompt(
+                quiz_state["language_mode"],
+                quiz_state["quiz_data"],
+                user_text
+            )
+            evaluation = chat_json(evaluation_prompt)
+
+            lines = []
+            lines.append(f"**Score:** {evaluation['score_obtained']}/{evaluation['score_total']}")
+            lines.append("")
+            lines.append(f"**Overall feedback:** {evaluation['summary']}")
+            lines.append("")
+            lines.append("**Question-wise evaluation:**")
+
+            for item in evaluation["results"]:
+                lines.append("")
+                lines.append(f"**Q:** {item['question']}")
+                lines.append(f"**Your answer:** {item['student_answer']}")
+                lines.append(f"**Result:** {item['result']}")
+                lines.append(f"**Feedback:** {item['feedback']}")
+
+            final_answer = "\n".join(lines)
+            history.append({"role": "assistant", "content": final_answer})
+
+            quiz_state = {
+                "active": False,
+                "topic": None,
+                "quiz_data": None,
+                "language_mode": language_mode
+            }
+
+            return history, quiz_state, ""
+
+        # -----------------------------------------
+        # NORMAL RETRIEVAL
+        # -----------------------------------------
+        records = search_hybrid(user_text, shortlist_k=20, final_k=3)
+        context = build_context(records)
+
+        # -----------------------------------------
+        # QUIZ GENERATION
+        # -----------------------------------------
+        if mode == "Quiz Me":
+            n_questions = choose_quiz_count(user_text, quiz_count_mode)
+            prompt = build_quiz_generation_prompt(language_mode, user_text, context, n_questions)
+            quiz_data = chat_json(prompt)
+
+            lines = []
+            lines.append(f"**{quiz_data.get('title', 'Quiz')}**")
+            lines.append("")
+            lines.append("Please answer the following questions in one message.")
+            lines.append("You can reply in numbered format, for example:")
+            lines.append("1. ...")
+            lines.append("2. ...")
+            lines.append("")
+            lines.append(f"**Total questions: {len(quiz_data['questions'])}**")
+            lines.append("")
+
+            for i, q in enumerate(quiz_data["questions"], start=1):
+                lines.append(f"**Q{i}.** {q['q']}")
+
+            if show_sources:
+                lines.append("\n---\n**Topic sources used to create the quiz:**")
+                lines.append(make_sources(records))
+
+            assistant_text = "\n".join(lines)
+            history.append({"role": "assistant", "content": assistant_text})
+
+            quiz_state = {
+                "active": True,
+                "topic": user_text,
+                "quiz_data": quiz_data,
+                "language_mode": language_mode
+            }
+
+            return history, quiz_state, ""
+
+        # -----------------------------------------
+        # OTHER MODES
+        # -----------------------------------------
+        prompt = build_tutor_prompt(mode, language_mode, user_text, context)
+        answer = chat_text(prompt)
+
+        if show_sources:
+            answer += "\n\n---\n**Sources used:**\n" + make_sources(records)
+
+        history.append({"role": "assistant", "content": answer})
+        return history, quiz_state, ""
+
+    except Exception as e:
+        history.append({"role": "assistant", "content": f"Error: {str(e)}"})
+        quiz_state["active"] = False
+        return history, quiz_state, ""
+
+
+def clear_all():
+    empty_quiz = {
+        "active": False,
+        "topic": None,
+        "quiz_data": None,
+        "language_mode": "Auto"
+    }
+    return [], empty_quiz, ""
+
+
+CSS = """
+body, .gradio-container {
+    background: #dcdcdc !important;
+    font-family: Arial, Helvetica, sans-serif !important;
+}
+footer { display: none !important; }
+.hero-card {
+    max-width: 860px;
+    margin: 18px auto 14px auto;
+    border-radius: 28px;
+    background: linear-gradient(180deg, #e8c7d4 0%, #a55ca2 48%, #2b0c46 100%);
+    padding: 22px 22px 18px 22px;
+}
+.hero-inner { text-align: center; }
+.hero-title {
+    color: white;
+    font-size: 34px;
+    font-weight: 800;
+    margin-top: 6px;
+}
+.hero-subtitle {
+    color: white;
+    opacity: 0.92;
+    font-size: 16px;
+    margin-top: 6px;
+}
+"""
+
+
+with gr.Blocks(css=CSS) as demo:
+    quiz_state = gr.State({
+        "active": False,
+        "topic": None,
+        "quiz_data": None,
+        "language_mode": "Auto"
+    })
+
+    gr.HTML(render_header())
+
+    with gr.Row():
+        mode = gr.Dropdown(
+            choices=["Explain", "Detailed", "Short Notes", "Quiz Me", "Flashcards", "Case-Based"],
+            value="Explain",
+            label="Tutor Mode"
+        )
+        language_mode = gr.Dropdown(
+            choices=["Auto", "English", "Spanish", "Bilingual"],
+            value="Auto",
+            label="Answer Language"
+        )
+
+    with gr.Row():
+        quiz_count_mode = gr.Dropdown(
+            choices=["Auto", "3", "5", "7"],
+            value="Auto",
+            label="Quiz Questions"
+        )
+        show_sources = gr.Checkbox(value=True, label="Show Sources")
+
+    gr.Markdown("""
+**How to use**
+- Choose a **Tutor Mode**
+- Then type a topic or question
+- For **Quiz Me**, type a topic such as: `cranial nerves`
+- The system will ask questions, and your **next message will be evaluated automatically**
+""")
+
+    chatbot = gr.Chatbot(height=520, type="messages")
+    msg = gr.Textbox(
+        placeholder="Ask a question or type a topic...",
+        lines=1,
+        show_label=False
+    )
+
+    with gr.Row():
+        send_btn = gr.Button("Send")
+        clear_btn = gr.Button("Clear Chat")
+
+    msg.submit(
+        answer_question,
+        inputs=[msg, chatbot, mode, language_mode, quiz_count_mode, show_sources, quiz_state],
+        outputs=[chatbot, quiz_state, msg]
+    )
+
+    send_btn.click(
+        answer_question,
+        inputs=[msg, chatbot, mode, language_mode, quiz_count_mode, show_sources, quiz_state],
+        outputs=[chatbot, quiz_state, msg]
+    )
+
+    clear_btn.click(
+        clear_all,
+        inputs=[],
+        outputs=[chatbot, quiz_state, msg]
+    )
+
+if __name__ == "__main__":
+    demo.launch()