Views
No views yet
pip install transformers
pip install flash-attn --no-build-isolation1import json
2import torch
3from transformers import AutoModelForCausalLM, AutoTokenizer
4
5model_id = "namespace-Pt/beacon-qwen-2-7b-instruct"
6
7tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
8model = AutoModelForCausalLM.from_pretrained(
9 model_id,
10 trust_remote_code=True,
11 torch_dtype=torch.bfloat16,
12 attn_implementation="flash_attention_2"
13)
14
15model = model.cuda().eval()
16
17with torch.no_grad():
18 # short context
19 messages = [{"role": "user", "content": "Tell me about yourself."}]
20 inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
21 outputs = model.generate(**inputs, max_new_tokens=50)
22 print(f"Input Length: {inputs['input_ids'].shape[1]}")
23 print(f"Output: {repr(tokenizer.decode(outputs[0], skip_special_tokens=True))}")
24
25 # reset memory before new generation task
26 model.memory.reset()
27
28 # long context
29 with open("infbench.json", encoding="utf-8") as f:
30 example = json.load(f)
31 messages = [{"role": "user", "content": example["context"]}]
32 inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to("cuda")
33 outputs = model.generate(**inputs, do_sample=False, top_p=1, temperature=1, max_new_tokens=20)[:, inputs["input_ids"].shape[1]:]
34 print("*"*20)
35 print(f"Input Length: {inputs['input_ids'].shape[1]}")
36 print(f"Answers: {example['answer']}")
37 print(f"Prediction: {tokenizer.decode(outputs[0], skip_special_tokens=True)}")This is a friendly reminder - the current text generation call will exceed the model's predefined maximum length (32768). Depending on the model, you may observe exceptions, performance degradation, or nothing at all. Just ignore it.