eduintel_generate.py

#!/usr/bin/env python3
import sys, os, pickle, re, json, warnings
import numpy as np
import requests
from pathlib import Path
from sentence_transformers import SentenceTransformer
import faiss
warnings.filterwarnings("ignore")

INDEX_DIR   = Path.home() / "leotsha_project/eduintel_index"
TRANS_MODEL = "Helsinki-NLP/opus-mt-en-nso"
EMBED_MODEL = "paraphrase-multilingual-MiniLM-L12-v2"

GRADE_FOLDERS = {
    "mind_the_gap","self_study","caps","atp",
    "2017_nov","2018_nov","2019_jun","2019_nov",
    "2020_nov","2021_jun","2021_nov","2022_jun","2022_nov",
    "2023_jun","2023_nov","2024_jun","2024_nov",
    "2025_jun","2025_nov","gr10_common","gr11_common",
}

GLOSSARY = {
    r"\bGrade\s+10\b": "mphato wa 10",
    r"\bGrade\s+11\b": "mphato wa 11",
    r"\bGrade\s+12\b": "mphato wa 12",
    r"\blearners\b":   "baithuti",
    r"\bteachers\b":   "barutiši",
    r"\bnovel\b":      "padi",
    r"\bpoems\b":      "disereto",
    r"\bpoem\b":       "sereto",
    r"\btheme\b":      "molaetša",
    r"\bexplain\b":    "hlaloša",
    r"\bteach\b":      "ruta",
    r"\bgrammar\b":    "puo",
    r"\bprescribed\b": "laetšwego",
    r"\bterm\s+1\b":   "kotara ya 1",
    r"\bterm\s+2\b":   "kotara ya 2",
    r"\bterm\s+3\b":   "kotara ya 3",
    r"\bHome\s+Language\b": "Leleme la Gae",
    r"\bSepedi\s+HL\b": "Sepedi Leleme la Gae",
    r"\bBadimo\s+Ba\s+Boletse\b": "Badimo Ba Boletše",
    r"\bancestors\b":    "badimo",
    r"\bancestor\b":     "modimo wa leloko",
    r"\bprescribed\s+novel\b": "padi ya go laetšwa",
    r"\bcharacter\b":    "moswantšho",
    r"\bplot\b":         "kgang ya padi",
    r"\bsetting\b":      "lefelo la padi",
    r"\bconflict\b":     "kgang",
    r"\bmoral\b":        "molaetša wa boitshwaro",
    r"\btradition\b":    "setšo",
    r"\bculture\b":      "setšo",
    r"\bidentity\b":     "boitsebišo",
    r"\bO\s+iphihletšeng\b": "O iphihletšeng",
    r"\bMontshepetsabosego\b": "Montshepetšabošego",
}

def apply_glossary(text):
    for pattern, replacement in GLOSSARY.items():
        text = re.sub(pattern, replacement, text, flags=re.IGNORECASE)
    return text

def translate(text):
    preprocessed = apply_glossary(text)
    # Count Sepedi terms already present after glossary
    sepedi_markers = ["mphato", "kotara", "baithuti", "barutiši",
                      "padi", "sereto", "disereto", "molaetša",
                      "hlaloša", "ruta", "puo", "kwešišo", "tekolo",
                      "Leleme", "Kreiti", "Beke", "CAPS"]
    sepedi_count = sum(1 for m in sepedi_markers if m.lower() in preprocessed.lower())

    # If 2+ Sepedi terms already present, skip Helsinki-NLP translation
    # The glossary has done enough — Helsinki-NLP will corrupt the good terms
    if sepedi_count >= 2:
        return preprocessed

    # Otherwise translate with Helsinki-NLP
    from transformers import MarianMTModel, MarianTokenizer
    tok = MarianTokenizer.from_pretrained(TRANS_MODEL)
    mdl = MarianMTModel.from_pretrained(TRANS_MODEL)
    inputs = tok([preprocessed], return_tensors="pt", padding=True)
    out = mdl.generate(**inputs, max_length=256)
    return tok.decode(out[0], skip_special_tokens=True)

def retrieve(query_nso, top_k=5):
    NOVEL_KEYWORDS = {
        "badimo ba boletše": "Badimo_Ba_Boletše.pdf",
        "badimo ba boletse": "Badimo_Ba_Boletše.pdf",
        "o iphihletšeng":    "O_iphihletšeng.pdf",
        "montshepetšabošego": "Montshepetšabošego.pdf",
        "ntlhomole mmutlwa": "Ntlhomole_Mmutlwa.pdf",
    }
    index  = faiss.read_index(str(INDEX_DIR / "eduintel.index"))
    chunks = pickle.load(open(INDEX_DIR / "chunks.pkl", "rb"))
    model  = SentenceTransformer(EMBED_MODEL)
    q_vec  = model.encode([query_nso], normalize_embeddings=True).astype(np.float32)
    scores, indices = index.search(q_vec, 80)

    query_lower = query_nso.lower()
    boost_file  = next(
        (fname for kw, fname in NOVEL_KEYWORDS.items() if kw in query_lower),
        None
    )

    results = []
    seen = {}
    for score, idx in zip(scores[0], indices[0]):
        if idx < 0 or idx >= len(chunks): continue
        chunk  = chunks[idx]
        folder = chunk.get("folder", "")
        fname  = chunk.get("file", "")
        if folder not in GRADE_FOLDERS: continue
        boosted = float(score) + (0.15 if boost_file and fname == boost_file else 0)
        if seen.get(fname, 0) >= 3: continue
        seen[fname] = seen.get(fname, 0) + 1
        results.append((boosted, chunk))

    results.sort(key=lambda x: -x[0])
    return results[:top_k]

def generate_answer(query_en, results):
    context_parts = []
    for i, (score, chunk) in enumerate(results, 1):
        folder = chunk.get("folder", "")
        fname  = chunk.get("file", "")
        text   = chunk.get("text", "").strip()
        context_parts.append(
            f"[Source {i} — {folder}/{fname} ({int(score*100)}% match)]\n{text[:600]}"
        )
    context = "\n\n".join(context_parts)

    system = """You are EduIntel, an AI assistant for Sepedi language teachers in South Africa (Grade 10-12).
You help Morutabana (teacher) users with CAPS curriculum questions.

CRITICAL RULES:
- ONLY use information explicitly present in the provided sources
- NEVER invent poem titles, novel titles, or specific content not in the sources
- Badimo Ba Boletše is a SEPEDI novel — never refer to it as Setswana
- All novels in the Sepedi HL curriculum are in SEPEDI, not Setswana or Zulu
- Answer in ENGLISH only — do not mix languages in your response
- Cite sources by their document name when making specific claims
- If the sources don't contain enough detail, say so clearly and suggest the teacher consult the CAPS document directly
- Be practical and structured but stay grounded in retrieved content"""

    user = f"Teacher question: {query_en}\n\nRelevant curriculum content:\n{context}\n\nProvide a helpful, practical answer."

    import requests as req
    key = os.environ.get("GROQ_KEY", "")
    r = req.post(
        "https://api.groq.com/openai/v1/chat/completions",
        headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
        json={
            "model":    "llama-3.1-8b-instant",
            "messages": [
                {"role": "system", "content": system},
                {"role": "user",   "content": user}
            ],
            "max_tokens": 800
        },
        timeout=30
    )
    data = r.json()
    if "choices" in data:
        return data["choices"][0]["message"]["content"]
    return f"Error: {data}"

def main():
    if len(sys.argv) < 2:
        print('Usage: python3 eduintel_generate.py "your question"')
        return
    query_en = " ".join(sys.argv[1:])
    print(f"\n{'━'*60}")
    print("  EduIntel — Synthesized Answer")
    print(f"{'━'*60}")
    print(f"\n  Question: {query_en}\n")
    print("  Translating...")
    query_nso = translate(query_en)
    print(f"  Sepedi: {query_nso}\n")
    print("  Retrieving relevant content...")
    results = retrieve(query_nso)
    print(f"  Found {len(results)} passages\n")
    print("  Generating answer...\n")
    answer = generate_answer(query_en, results)
    print(f"{'─'*60}")
    print(answer)
    print(f"{'━'*60}\n")

if __name__ == "__main__":
    main()