import sysfrom pathlib import Pathimport torchfrom huggingface_hub import snapshot_downloadfrom transformers import AutoModelForCausalLM, AutoTokenizer model_dir = Path(snapshot_download("AI45Research/ReSI-DeepSeek-V4-Flash"))sys.path.insert(0, str(model_dir / "encoding"))from encoding_dsv4 import encode_messages tokenizer = AutoTokenizer.from_pretrained(model_dir)model = AutoModelForCausalLM.from_pretrained( model_dir, dtype=torch.bfloat16, device_map="balanced", attn_implementation="eager",).eval()text = encode_messages( [{"role": "user", "content": "Explain why the sky is blue in three sentences."}], thinking_mode="chat",)inputs = tokenizer(text, add_special_tokens=False, return_tensors="pt").to( model.get_input_embeddings().weight.device)with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): outputs = model.generate( **inputs, max_new_tokens=8192, do_sample=True, temperature=1.0, top_p=1.0, top_k=0, )print(tokenizer.decode(outputs[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True))