Spaces:
Runtime error
Runtime error
| import os | |
| import json | |
| import torch | |
| import torch.nn as nn | |
| import sentencepiece as spm | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| import spaces | |
| # --- 1. CONFIGURATION & MODEL DOWNLOAD --- | |
| REPO_ID = "yisberh/grade8-math-custom-transformer" | |
| # Fetch preloaded cache files instantly from local disk memory | |
| config_path = hf_hub_download(repo_id=REPO_ID, filename="config.json", local_files_only=True) | |
| model_path = hf_hub_download(repo_id=REPO_ID, filename="model.pt", local_files_only=True) | |
| spm_path = hf_hub_download(repo_id=REPO_ID, filename="unit1_test.model", local_files_only=True) | |
| with open(config_path, "r", encoding="utf-8") as f: | |
| config = json.load(f) | |
| sp = spm.SentencePieceProcessor(model_file=spm_path) | |
| # --- 2. MODEL DEFINITION (Matching your exact architecture) --- | |
| class Model(nn.Module): | |
| def __init__(self, vocab): | |
| super().__init__() | |
| self.d = 256 # 256 embedding dimensions | |
| self.emb = nn.Embedding(vocab, self.d) | |
| self.pos = nn.Embedding(512, self.d) | |
| layer = nn.TransformerEncoderLayer( | |
| d_model=self.d, | |
| nhead=8, # 8 attention heads | |
| dim_feedforward=1024, # 1024 FFN dimension | |
| batch_first=True, | |
| dropout=0.0 # 0.0 dropout for speed during inference | |
| ) | |
| self.tr = nn.TransformerEncoder(layer, 6) # 6 Layers deep | |
| self.head = nn.Linear(self.d, vocab) | |
| def forward(self, x): | |
| n = x.size(1) | |
| pos = torch.arange(n, device=x.device) | |
| x = self.emb(x) + self.pos(pos) | |
| mask = torch.triu( | |
| torch.full((n, n), float("-inf"), device=x.device), | |
| diagonal=1 | |
| ) | |
| return self.head(self.tr(x, mask=mask)) | |
| # Initialize model container framework on CPU first | |
| model = Model(sp.vocab_size()) | |
| # Strip any potential DataParallel state_dict prefix formatting cleanly | |
| state_dict = torch.load(model_path, map_location="cpu") | |
| if any(k.startswith("module.") for k in state_dict.keys()): | |
| state_dict = {k.replace("module.", ""): v for k in state_dict.items()} | |
| model.load_state_dict(state_dict) | |
| model.eval() | |
| # --- 3. INFERENCE LOOP ROUTED VIA ZEROGPU --- | |
| # Mandated to prevent "No @spaces.GPU function detected" | |
| def generate_response(question): | |
| # Protect against early app validation/mock checks passing empty inputs | |
| if not question or not isinstance(question, str) or not question.strip(): | |
| return "Please ask a question about Grade 8 Mathematics Unit 1." | |
| # Identify device assigned dynamically by the Hugging Face manager | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model.to(device) | |
| # Encode instructions using your native data pipeline format | |
| ids = [sp.bos_id()] + sp.encode_as_ids(f"[INST] {question.strip()} [/INST]") | |
| x = torch.tensor([ids], device=device) | |
| max_new_tokens = 256 # Higher generation limits to support extended answers | |
| for _ in range(max_new_tokens): | |
| if x.size(1) >= 512: # Strict constraint boundary matching sequence history size | |
| break | |
| logits = model(x) | |
| nxt = logits[:, -1].argmax(-1).item() | |
| nxt_tensor = torch.tensor([[nxt]], device=device) | |
| x = torch.cat([x, nxt_tensor], dim=1) | |
| if nxt == sp.eos_id(): | |
| break | |
| text = sp.decode(x[0].tolist()) | |
| # Split using the explicit token tags matching your SFT layout formatting | |
| if "[/INST]" in text: | |
| return text.split("[/INST]")[-1].strip() | |
| return text | |
| # --- 4. GRADIO INTERFACE LAYOUT CONFIGURATION --- | |
| demo = gr.Interface( | |
| fn=generate_response, | |
| inputs=gr.Textbox( | |
| lines=3, | |
| placeholder="Enter your Grade 8 Mathematics question here... (Supports English & Amharic)", | |
| label="Math Question" | |
| ), | |
| outputs=gr.Textbox(label="Model Response"), | |
| title="Grade 8 Mathematics AI Assistant (5M Params)", | |
| description="Custom transformer model tracking Grade 8 Math Unit 1 roadmap objectives." | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(server_name="0.0.0.0", server_port=7860) | |