app.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import re
import hashlib
import json
import torch

# ============================================================
# DEVICE
# ============================================================

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ============================================================
# MODELS
# ============================================================

SIM_MODEL_NAME = "cross-encoder/stsb-distilroberta-base"
NLI_MODEL_NAME = "cross-encoder/nli-deberta-v3-xsmall"
LLM_NAME = "google/flan-t5-base"

print("Loading similarity + NLI models...")
sim_model = CrossEncoder(SIM_MODEL_NAME, device=DEVICE)
nli_model = CrossEncoder(NLI_MODEL_NAME, device=DEVICE)

print("Loading LLM for atomic fact extraction...")
llm_tokenizer = AutoTokenizer.from_pretrained(LLM_NAME)
llm_model = AutoModelForSeq2SeqLM.from_pretrained(LLM_NAME).to(DEVICE)

print("✅ All models loaded")

# ============================================================
# CONFIGURATION
# ============================================================
SIM_THRESHOLD_REQUIRED = 0.55

CONTRADICTION_THRESHOLD = 0.70
SCHEMA_CACHE = {}

# ============================================================
# UTILITIES
# ============================================================

def split_sentences(text):
    return re.split(r'(?<=[.!?])\s+', text.strip())

def hash_key(kb, question):
    return hashlib.sha256((kb + question).encode()).hexdigest()









def decompose_answer(answer):
    """Split answer into atomic claims."""
    parts = re.split(r'\b(?:and|because|before|after|while|then|so)\b', answer)
    return [p.strip() for p in parts if p.strip()]

# ============================================================
# LLM FACT EXTRACTION
# ============================================================
def generate_atomic_facts(kb, question):
    """
    Ask LLM to extract 1-5 atomic facts from KB that directly answer the question.
    Returns JSON: {"facts": [ ... ]}
    """
    prompt = f"""
Extract atomic facts that directly answer the question.










Knowledge Base:
{kb}
Question:
{question}

RULES:
- Return 1-5 short factual statements that directly answer the question.
- Output strictly in JSON format: {{"facts": ["fact1", "fact2", ...]}}
- Do not include unrelated events or explanations.
- Each fact should be self-contained.
"""

    inputs = llm_tokenizer(prompt, return_tensors="pt", truncation=True).to(DEVICE)
    outputs = llm_model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        temperature=0.0
    )

    raw = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    try:
        data = json.loads(raw)
        facts = data.get("facts", [])
    except:
        # fallback: parse line by line if JSON fails
        facts = [line.strip("-• ").strip() for line in raw.split("\n") if len(line.strip()) > 3]
    
    return {

        "required_concepts": facts,
        "raw_llm_output": raw
    }
# ============================================================
# CORE EVALUATION
# ============================================================

def evaluate_answer(answer, question, kb):
    logs = {"inputs": {"question": question, "answer": answer, "kb_length": len(kb)}}
    






    key = hash_key(kb, question)

    if key not in SCHEMA_CACHE:
        schema = generate_atomic_facts(kb, question)















        SCHEMA_CACHE[key] = schema
    
    schema = SCHEMA_CACHE[key]
    logs["schema"] = schema
    
    claims = decompose_answer(answer)
    logs["claims"] = claims
    
    # ---------------- COVERAGE ----------------
    coverage = []
    covered_all = True

    for concept in schema["required_concepts"]:
        if claims:
            scores = sim_model.predict([(concept, c) for c in claims])
            best = float(scores.max())
            ok = best >= SIM_THRESHOLD_REQUIRED
        else:
            best = 0.0
            ok = False
        coverage.append({
            "concept": concept,
            "similarity": round(best, 3),
            "covered": ok
        })

        if not ok:
            covered_all = False

    logs["coverage"] = coverage
    
    # ---------------- CONTRADICTIONS ----------------
    contradictions = []
    kb_sents = split_sentences(kb)
    for claim in claims:
        for sent in kb_sents:
            probs = softmax_logits(nli_model.predict([(sent, claim)]))
            if probs[0] > CONTRADICTION_THRESHOLD:
                contradictions.append({
                    "claim": claim,
                    "sentence": sent,
                    "confidence": round(probs[0] * 100, 1)
                })

    logs["contradictions"] = contradictions
    
    # ---------------- FINAL VERDICT ----------------
    if contradictions:
        verdict = "❌ INCORRECT (Contradiction)"
    elif covered_all:
        verdict = "✅ CORRECT"
    else:
        verdict = "⚠️ PARTIALLY CORRECT"
    
    logs["final_verdict"] = verdict
    return verdict, logs

# ============================================================
# GRADIO UI
# ============================================================

def run(answer, question, kb):
    return evaluate_answer(answer, question, kb)

with gr.Blocks(title="Competitive Exam Answer Checker") as demo:
    gr.Markdown("## 🧠 Competitive Exam Answer Checker (Robust General Version)")
    
    kb = gr.Textbox(label="Knowledge Base", lines=10)
    question = gr.Textbox(label="Question")
    answer = gr.Textbox(label="Student Answer")
    
    verdict = gr.Textbox(label="Verdict")
    debug = gr.JSON(label="Debug Logs")
    
    btn = gr.Button("Evaluate")
    btn.click(run, [answer, question, kb], [verdict, debug])