Buckets:
| """ | |
| inference.py — plug-and-play inference for AxionLab-official/SupraElegans-500k | |
| Usage: | |
| python inference.py --prompt "Once upon a time" | |
| python inference.py --prompt "The little robot" --max_new_tokens 150 --temperature 0.7 | |
| python inference.py --interactive | |
| Requirements: | |
| pip install torch transformers | |
| The model repo is loaded with trust_remote_code=True because SupraElegans is | |
| not a Transformer and its architecture (modeling_supraelegans.py, shipped in | |
| the repo) is not part of the transformers library. | |
| """ | |
| import argparse | |
| import torch | |
| from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedTokenizerFast | |
| from configuration_supraelegans import SupraElegansConfig | |
| from modeling_supraelegans import SupraElegansForCausalLM | |
| MODEL_ID = "SupraLabs/SupraElegans-500k" | |
| # Registers the local model/config classes with the Auto* factories so that | |
| # `AutoModelForCausalLM.from_pretrained` resolves `model_type: "supraelegans"` | |
| # without needing an `auto_map` entry in the repo's config.json, and without | |
| # depending on `trust_remote_code` fetching/executing a remote file at all. | |
| # This is what makes loading robust regardless of which tool the checkpoint | |
| # was originally saved with (some don't write `auto_map`). | |
| AutoConfig.register("supraelegans", SupraElegansConfig) | |
| AutoModelForCausalLM.register(SupraElegansConfig, SupraElegansForCausalLM) | |
| def load_model(model_id=MODEL_ID, device=None): | |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") | |
| # Loaded via PreTrainedTokenizerFast directly rather than AutoTokenizer. | |
| # Some environments the tokenizer was saved from write a | |
| # "tokenizer_class" value into tokenizer_config.json that older | |
| # `transformers` installs don't recognize, which makes AutoTokenizer | |
| # raise instead of falling back. The tokenizer is a plain fast BPE | |
| # tokenizer, so loading the concrete class sidesteps that lookup. | |
| tokenizer = PreTrainedTokenizerFast.from_pretrained(model_id) | |
| model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True) | |
| model.to(device) | |
| model.eval() | |
| return model, tokenizer, device | |
| def generate(model, tokenizer, device, prompt, max_new_tokens=100, | |
| temperature=0.8, top_k=50, top_p=0.9): | |
| ids = tokenizer.encode(prompt) | |
| bos = tokenizer.bos_token_id | |
| input_ids = torch.tensor([[bos] + ids if bos is not None else ids], device=device) | |
| output_ids, _ = model.generate( | |
| input_ids, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature, | |
| top_k=top_k, | |
| top_p=top_p, | |
| ) | |
| return tokenizer.decode(output_ids[0].tolist(), skip_special_tokens=True) | |
| def run_single(args): | |
| model, tokenizer, device = load_model(args.model_id, args.device) | |
| text = generate( | |
| model, tokenizer, device, args.prompt, | |
| max_new_tokens=args.max_new_tokens, | |
| temperature=args.temperature, | |
| top_k=args.top_k, | |
| top_p=args.top_p, | |
| ) | |
| print(text) | |
| def run_interactive(args): | |
| model, tokenizer, device = load_model(args.model_id, args.device) | |
| print(f"SupraElegans-500K loaded on {device}. Ctrl+C or empty input to exit.\n") | |
| while True: | |
| try: | |
| prompt = input("> ").strip() | |
| except (EOFError, KeyboardInterrupt): | |
| break | |
| if not prompt: | |
| break | |
| text = generate( | |
| model, tokenizer, device, prompt, | |
| max_new_tokens=args.max_new_tokens, | |
| temperature=args.temperature, | |
| top_k=args.top_k, | |
| top_p=args.top_p, | |
| ) | |
| print(text, "\n") | |
| def main(): | |
| parser = argparse.ArgumentParser(description="SupraElegans-500K inference") | |
| parser.add_argument("--model_id", type=str, default=MODEL_ID, | |
| help="Hugging Face repo id or local path") | |
| parser.add_argument("--prompt", type=str, default=None, | |
| help="Prompt text (ignored in --interactive mode)") | |
| parser.add_argument("--max_new_tokens", type=int, default=100) | |
| parser.add_argument("--temperature", type=float, default=0.8) | |
| parser.add_argument("--top_k", type=int, default=50) | |
| parser.add_argument("--top_p", type=float, default=0.9) | |
| parser.add_argument("--device", type=str, default=None, | |
| help="cuda, cpu, or leave unset to auto-detect") | |
| parser.add_argument("--interactive", action="store_true", | |
| help="Drop into a REPL instead of a single generation") | |
| args = parser.parse_args() | |
| if args.interactive: | |
| run_interactive(args) | |
| else: | |
| if not args.prompt: | |
| parser.error("--prompt is required unless --interactive is set") | |
| run_single(args) | |
| if __name__ == "__main__": | |
| main() | |
Xet Storage Details
- Size:
- 4.79 kB
- Xet hash:
- ac3320181ad8f3a9d21bda52a6caf5f7f96be61c553f77d3f0f2f92feb832d80
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.