Update app.py

app.py

@@ -1,31 +1,40 @@
 """
-Estella Explainer — HuggingFace Spaces / Lightning AI Demo
-A math reading coach for grades 3–8, based on Usable Math (usablemath.org)
-
-Developed by Sai Gattupalli, PhD
-College of Education, University of Massachusetts Amherst
-Principal Scientist, Society and AI Research Group
-
-Model: sgattup/EstellaExplainerLLM (Gemma 4 E4B fine-tuned)
+Estella Explainer — Hugging Face Space
+A math reading coach for grades 3–8.
 """
 
+import os
+
 import gradio as gr
 import torch
+from peft import PeftModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_ID = "sgattup/EstellaExplainerLLM"
 
-SYSTEM_PROMPT = """You are Estella Explainer, a math reading and language coach for young learners in grades 3 through 8. Your motto is: "My job is to explain math questions clearly so you know what you are supposed to do to solve the problem."
+ADAPTER_ID = "sgattup/EstellaExplainerLLM"
+BASE_MODEL_ID = "unsloth/gemma-4-E4B-it"
+
+HF_TOKEN = os.getenv("HF_TOKEN")
+HUB_ARGS = {"token": HF_TOKEN} if HF_TOKEN else {}
 
-Your rules:
+
+SYSTEM_PROMPT = """You are Estella Explainer, a math reading and language coach for young learners in grades 3 through 8.
+
+Your motto is:
+"My job is to explain math questions clearly so you know what you are supposed to do to solve the problem."
+
+Rules:
 - Give only ONE hint per problem.
 - Do NOT solve the problem or show the answer.
-- Do NOT show calculation steps or compute anything.
-- Help the student understand what the problem is ASKING.
-- Use very simple words and short sentences (Flesch Reading Ease 90-100).
+- Do NOT show calculation steps.
+- Do NOT compute anything.
+- Help the student understand what the problem is asking.
+- Use very simple words and short sentences.
 - Identify what is known and what is being asked.
 - Use phrases like "We know...", "We are looking for...", or "Think about...".
-- Be friendly and encouraging."""
+- Be friendly and encouraging.
+"""
+
 
 EXAMPLES = [
     "A rectangle has a length of 12 cm and a width of 5 cm. What is the area?",
@@ -38,33 +47,85 @@ EXAMPLES = [
     "The perimeter of a square is 36 inches. What is the length of one side?",
 ]
 
-print(f"Loading {MODEL_ID}...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    torch_dtype=torch.float16,
-    device_map="auto"
+
+print(f"Loading base model: {BASE_MODEL_ID}")
+
+if torch.cuda.is_available():
+    model_dtype = torch.float16
+    device_map = "auto"
+    print("GPU detected. Loading model in float16.")
+else:
+    model_dtype = torch.float32
+    device_map = None
+    print("No GPU detected. CPU loading may be slow or require more memory.")
+
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_ID,
+    torch_dtype=model_dtype,
+    device_map=device_map,
+    low_cpu_mem_usage=True,
+    **HUB_ARGS,
+)
+
+
+print(f"Loading Estella adapter: {ADAPTER_ID}")
+
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER_ID,
+    **HUB_ARGS,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    ADAPTER_ID,
+    **HUB_ARGS,
 )
+
 model.eval()
+
 print("Estella is ready.")
 
 
 def ask_estella(message: str, history: list) -> str:
-    """Chat function — history is kept for UI display only; Estella gives one hint per problem."""
-    if not message.strip():
-        return "Please type a math word problem and I'll help explain what it's asking!"
+    """Return one reading hint without solving the math problem."""
+
+    if not message or not message.strip():
+        return "Please type a math word problem and I’ll help explain what it is asking!"
 
     messages = [
-        {"role": "user", "content": SYSTEM_PROMPT + f"\n\n### Math Problem:\n{message.strip()}\n\n### Estella's Hint:"}
+        {
+            "role": "system",
+            "content": SYSTEM_PROMPT,
+        },
+        {
+            "role": "user",
+            "content": (
+                "### Math Problem:\n"
+                f"{message.strip()}\n\n"
+                "### Estella's Hint:"
+            ),
+        },
     ]
-    text = tokenizer.apply_chat_template(
+
+    prompt = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
-        tokenize=False
+        tokenize=False,
+    )
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
     )
-    inputs = tokenizer(text=[text], return_tensors="pt")
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    input_len = inputs["input_ids"].shape[1]
+
+    input_device = next(model.parameters()).device
+    inputs = {
+        key: value.to(input_device)
+        for key, value in inputs.items()
+    }
+
+    input_length = inputs["input_ids"].shape[1]
 
     with torch.no_grad():
         outputs = model.generate(
@@ -73,24 +134,37 @@ def ask_estella(message: str, history: list) -> str:
             temperature=0.7,
             do_sample=True,
             repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id,
         )
 
-    hint = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
+    generated_tokens = outputs[0][input_length:]
+
+    hint = tokenizer.decode(
+        generated_tokens,
+        skip_special_tokens=True,
+    ).strip()
+
     return hint
 
 
-# Gradio ChatInterface
 demo = gr.ChatInterface(
     fn=ask_estella,
     title="📚 Estella Explainer — Math Reading Coach",
-    description="""**"My job is to explain math questions clearly so you know what you are supposed to do to solve the problem."**
-
-Type any math word problem (grades 3–8) and Estella will help you understand what it's asking — without solving it.
-
-*Developed by [Sai Gattupalli, PhD](https://huggingface.co/sgattup) · UMass Amherst College of Education · [Society and AI Research Group](https://societyandai.org) · Based on [Usable Math](https://usablemath.org)*""",
+    description=(
+        '**"My job is to explain math questions clearly so you know '
+        'what you are supposed to do to solve the problem."**\n\n'
+        "Type any math word problem and Estella will help you understand "
+        "what it is asking—without solving it.\n\n"
+        "Developed by [Sai Gattupalli, PhD](https://huggingface.co/sgattup) · "
+        "[UMass Amherst College of Education](https://www.umass.edu/education/) · "
+        "[Society and AI Research Group](https://societyandai.org) · "
+        "[Usable Math](https://usablemath.org)"
+    ),
     examples=EXAMPLES,
     cache_examples=False,
     type="messages",
 )
 
-demo.launch()
+
+if __name__ == "__main__":
+    demo.launch()