Spaces:
Sleeping
Sleeping
| """ | |
| Estella Explainer — streaming Gradio chat interface | |
| This version keeps Gradio's default appearance and adds the two important | |
| interaction patterns used by the DeepSeek demo: token streaming and a queue. | |
| """ | |
| import os | |
| from threading import Thread | |
| from typing import Iterator | |
| import gradio as gr | |
| import spaces | |
| import torch | |
| from peft import PeftModel | |
| from transformers import ( | |
| AutoModelForCausalLM, | |
| AutoTokenizer, | |
| TextIteratorStreamer, | |
| ) | |
| 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 {} | |
| SYSTEM_PROMPT = """You are Estella Explainer, a math reading and language coach for young learners in grades 3 through 8. | |
| Estella helps students understand what a math problem is asking before they begin solving it. | |
| Rules: | |
| - Give only ONE hint per problem. | |
| - Do NOT solve the problem or show the answer. | |
| - 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. | |
| """ | |
| EXAMPLES = [ | |
| "A rectangle has a length of 12 cm and a width of 5 cm. What is the area?", | |
| "Sam had 3/4 of a pizza. He ate 1/4. How much is left?", | |
| "A store sells apples for $0.75 each. Jake buys 6 apples. How much does he spend?", | |
| "Round 4,567 to the nearest hundred.", | |
| "There are 48 students put into equal groups of 6. How many groups are there?", | |
| "A bag has 5 red marbles and 3 blue marbles. What fraction of the marbles are red?", | |
| "A train leaves at 10:15 AM and arrives at 1:45 PM. How long is the trip?", | |
| "The perimeter of a square is 36 inches. What is the length of one side?", | |
| ] | |
| print(f"Loading base model: {BASE_MODEL_ID}") | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL_ID, | |
| dtype=torch.float16, | |
| low_cpu_mem_usage=True, | |
| **HUB_ARGS, | |
| ) | |
| print(f"Loading adapter on CPU: {ADAPTER_ID}") | |
| model = PeftModel.from_pretrained( | |
| base_model, | |
| ADAPTER_ID, | |
| torch_device="cpu", | |
| **HUB_ARGS, | |
| ) | |
| print("Moving combined model to CUDA.") | |
| model = model.to("cuda") | |
| model.eval() | |
| tokenizer = AutoTokenizer.from_pretrained( | |
| ADAPTER_ID, | |
| **HUB_ARGS, | |
| ) | |
| print("Estella is ready.") | |
| def ask_estella(message: str, history: list) -> Iterator[str]: | |
| """Stream one reading hint without solving the math problem.""" | |
| if not message or not message.strip(): | |
| yield ( | |
| "Please type a math word problem and I’ll help explain " | |
| "what it is asking!" | |
| ) | |
| return | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": SYSTEM_PROMPT, | |
| }, | |
| { | |
| "role": "user", | |
| "content": ( | |
| "Math Problem:\n" | |
| f"{message.strip()}\n\n" | |
| "Estella's Hint:" | |
| ), | |
| }, | |
| ] | |
| prompt = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| tokenize=False, | |
| ) | |
| inputs = tokenizer( | |
| prompt, | |
| return_tensors="pt", | |
| ) | |
| input_device = next(model.parameters()).device | |
| inputs = { | |
| key: value.to(input_device) | |
| for key, value in inputs.items() | |
| } | |
| streamer = TextIteratorStreamer( | |
| tokenizer, | |
| timeout=60.0, | |
| skip_prompt=True, | |
| skip_special_tokens=True, | |
| ) | |
| generation_kwargs = { | |
| **inputs, | |
| "streamer": streamer, | |
| "max_new_tokens": 150, | |
| "temperature": 0.7, | |
| "do_sample": True, | |
| "repetition_penalty": 1.1, | |
| "pad_token_id": tokenizer.eos_token_id, | |
| } | |
| generation_thread = Thread( | |
| target=model.generate, | |
| kwargs=generation_kwargs, | |
| ) | |
| generation_thread.start() | |
| partial_hint = "" | |
| for new_text in streamer: | |
| partial_hint += new_text | |
| yield partial_hint.strip() | |
| demo = gr.ChatInterface( | |
| fn=ask_estella, | |
| title="Estella Explainer Math Reading Coach", | |
| description=( | |
| "Estella helps students understand what a math problem is asking " | |
| "before they begin solving it.\n\n" | |
| "Enter a math question below for one clear, friendly hint. " | |
| "Estella will not give away the answer.\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, | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=20).launch() | |