import torchfrom tokenizers import pre_tokenizersfrom transformers import AutoModelForCausalLM, AutoTokenizer
REPO_ID = "SupraLabs/Supra2-Medium-Instruct"device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"[*] Loading model and tokenizer from Hub: {REPO_ID}...")
# 1. Load Tokenizer & configure ByteLevel pre-tokenizertokenizer = AutoTokenizer.from_pretrained(REPO_ID, trust_remote_code=True)tokenizer._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
# 2. Load Modelmodel = AutoModelForCausalLM.from_pretrained( REPO_ID, torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, device_map="auto" if torch.cuda.is_available() else None, trust_remote_code=True)model.eval()
# 3. Prepare Chat Promptmessages = [ {"role": "user", "content": "What is AI?"}]
prompt_text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt_text, return_tensors="pt").to(device)
print("[*] Generating response...\n")
# 4. Generateim_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")eos_ids = [im_end_id, tokenizer.eos_token_id] if im_end_id is not None else tokenizer.eos_token_id
with torch.no_grad(): output_ids = model.generate( **inputs, max_new_tokens=256, do_sample=True, temperature=0.2, top_p=0.85, top_k=25, no_repeat_ngram_size=3, pad_token_id=tokenizer.pad_token_id, eos_token_id=eos_ids )
# 5. Extract & Decode Responsegenerated_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]response = tokenizer.decode(generated_tokens, skip_special_tokens=True)
print("--- Output ---")print(response.strip())