Spaces:
Sleeping
Sleeping
| import torch | |
| import sentencepiece as spm | |
| from karModel import DictionaryTransformer, device | |
| # ------------------------- | |
| # 1. Load SentencePiece tokenizer | |
| # ------------------------- | |
| sp = spm.SentencePieceProcessor() | |
| sp.load("cam/dict_spm.model") | |
| print("✅ Loaded SentencePiece tokenizer") | |
| # Check if eos exists in vocab | |
| eos_id = sp.piece_to_id("<eos>") | |
| if eos_id == -1: | |
| print("⚠️ No <eos> token in vocab — will remove by string replace") | |
| def encode(s): | |
| return sp.encode(s, out_type=int) | |
| def decode(ids): | |
| # Drop EOS ids if tokenizer has them | |
| if eos_id != -1: | |
| ids = [i for i in ids if i != eos_id] | |
| text = sp.decode(ids) | |
| # Trim at the first <eos> if it somehow appears in raw text | |
| if "<eos>" in text: | |
| text = text.split("<eos>")[0] | |
| return text.strip() | |
| # ------------------------- | |
| # 2. Load trained model | |
| # ------------------------- | |
| vocab_size = sp.get_piece_size() | |
| model = DictionaryTransformer(vocab_size).to(device) | |
| model.load_state_dict(torch.load("cam/dictionary_transformer.pt", map_location=device)) | |
| model.eval() | |
| print("✅ Loaded trained DictionaryTransformer") | |
| # ------------------------- | |
| # 3. Prompt + Generation | |
| # ------------------------- | |
| def generate_definition(prompt, max_new_tokens=50, top_k=40): | |
| context = torch.tensor([encode(prompt)], dtype=torch.long, device=device) | |
| for _ in range(max_new_tokens): | |
| idx_cond = context[:, -model.position_embedding_table.num_embeddings:] | |
| logits, _ = model(idx_cond) | |
| logits = logits[:, -1, :] | |
| probs = torch.softmax(logits, dim=-1) | |
| # --- Top-k filtering --- | |
| if top_k is not None: | |
| values, indices = torch.topk(probs, k=top_k) | |
| probs = torch.zeros_like(probs).scatter_(1, indices, values) | |
| probs = probs / probs.sum(dim=-1, keepdim=True) | |
| next_id = torch.multinomial(probs, num_samples=1) | |
| context = torch.cat((context, next_id), dim=1) | |
| if eos_id != -1 and next_id.item() == eos_id: | |
| break | |
| return decode(context[0].tolist()) | |
| if __name__ == "__main__": | |
| while True: | |
| word = input("Enter a word (or 'quit'): ").strip() | |
| if word.lower() == "quit": | |
| break | |
| if not word.endswith(":"): | |
| word += ":" | |
| print("\n📖", generate_definition(word), "\n") | |