| """ |
| rewrite.py -- self-contained loader/inference for Escarda-Rewrite, an 86M prompt |
| optimizer. Given a raw user question it rewrites it into a detailed, role-grounded |
| instruction prompt for a downstream LLM (persona declaration + task framing + |
| clarifying questions). Pairs naturally with Escarda-86M-Identity as the downstream. |
| |
| Usage: |
| from rewrite import EscardaRewrite |
| rw = EscardaRewrite(".") |
| print(rw.rewrite("How do I lower my cholesterol?")) |
| # "As a Senior Medical Expert, your task is to provide clear guidance on ..." |
| |
| Files required: model.safetensors, config.json, tokenizer.json, tokenizer_config.json, |
| special_tokens_map.json, model_v2.py, config.py, spike_tokenizer.py |
| """ |
| import os, sys, torch |
|
|
| _HERE = os.path.dirname(os.path.abspath(__file__)) |
| if _HERE not in sys.path: |
| sys.path.insert(0, _HERE) |
| from model_v2 import SpikeWhaleLM |
| from spike_tokenizer import SpikeTokenizer |
|
|
| IM_START, IM_END = "<|im_start|>", "<|im_end|>" |
|
|
|
|
| class EscardaRewrite: |
| def __init__(self, path=_HERE, device=None, dtype=torch.float32): |
| self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| self.tok = SpikeTokenizer(os.path.join(path, "tokenizer.json")) |
| self.model = SpikeWhaleLM.from_pretrained(path).to(self.device).to(dtype).eval() |
| end = self.tok.encode(IM_END) |
| self.end_id = end[-1] if end else None |
| self.eos = getattr(self.tok, "eos_token_id", None) |
|
|
| @torch.no_grad() |
| def rewrite(self, query, max_new=120): |
| """Greedy decode -- the optimized prompt is a deterministic rewrite.""" |
| prompt = f"{IM_START}user\n{query}{IM_END}\n{IM_START}assistant\n" |
| ids = torch.tensor([self.tok.encode(prompt)], device=self.device) |
| start = ids.shape[1] |
| for _ in range(max_new): |
| nxt = int(self.model(ids).logits[0, -1].argmax()) |
| ids = torch.cat([ids, torch.tensor([[nxt]], device=self.device)], 1) |
| if nxt == self.end_id or (self.eos is not None and nxt == self.eos): |
| break |
| return self.tok.decode(ids[0, start:].tolist()).strip() |
|
|
|
|
| if __name__ == "__main__": |
| rw = EscardaRewrite(_HERE, device="cpu") |
| for q in ["How do I lower my cholesterol?", "How do I grow tomatoes?", |
| "What are my rights if I am fired?"]: |
| print(f"\nUSER: {q}\nREWRITE: {rw.rewrite(q)}") |
|
|
