Views
No views yet
1from transformers import AutoTokenizer, Qwen3ForTokenClassification, AttentionInterface
2from typing import Optional
3
4def register_fa_attention():
5 from flash_attn import flash_attn_func, flash_attn_varlen_func
6
7 def custom_attention_forward(
8 module: AttentionInterface,
9 query: torch.Tensor,
10 key: torch.Tensor,
11 value: torch.Tensor,
12 attention_mask: Optional[torch.Tensor] = None,
13 **kwargs,
14 ):
15 cu_seqlens_q = kwargs.get("cu_seqlens_q", None)
16 cu_seqlens_k = kwargs.get("cu_seqlens_k", None)
17 max_seqlen_q = kwargs.get("max_seqlen_q", None)
18 max_seqlen_k = kwargs.get("max_seqlen_k", None)
19 # permute query, key, value to (batch, seq_len, n_heads, head_dim)
20 query_permute = query.permute(0, 2, 1, 3)
21 key_permute = key.permute(0, 2, 1, 3)
22 value_permute = value.permute(0, 2, 1, 3)
23
24 if cu_seqlens_q is not None and cu_seqlens_k is not None:
25 attn_output = flash_attn_varlen_func(
26 q=query_permute.squeeze(0),
27 k=key_permute.squeeze(0),
28 v=value_permute.squeeze(0),
29 cu_seqlens_q=cu_seqlens_q,
30 cu_seqlens_k=cu_seqlens_k,
31 max_seqlen_q=max_seqlen_q,
32 max_seqlen_k=max_seqlen_k,
33 causal=False,
34 )
35 else:
36 attn_output = flash_attn_func(
37 query_permute, key_permute, value_permute,
38 causal=False,
39 )
40 return attn_output , None
41
42 AttentionInterface.register("fa_noncausal", custom_attention_forward)
43
44# Register custom non-causal FA (Feel free to use FA2/FA3), required GPU
45register_fa_attention()
46
47def tokenize_sentence_to_word(sentence:str ):
48 tokens = []
49 chinese_char_pattern = re.compile(r'[\u4e00-\u9fff]')
50 # Split text by spaces first
51 parts = sentence.split()
52 for part in parts:
53 if chinese_char_pattern.search(part):
54 # Character-level tokenization for Chinese
55 tokens.extend(list(part))
56 else:
57 # Word-level tokenization for other languages
58 tokens.append(part)
59 return tokens
60
61tokenizer = AutoTokenizer.from_pretrained("Scicom-intl/multilingual-dynamic-entity-decoder")
62model = Qwen3ForTokenClassification.from_pretrained(
63 "Scicom-intl/multilingual-dynamic-entity-decoder",
64 attn_implementation="fa_noncausal",
65 dtype=torch.bfloat16,
66 device_map={"":"cuda:0"}
67)
68
69word_token = tokenize_sentence_to_word("Hi, my name is Alex and I'm from Perlis")
70token = tokenizer(
71 word_token,
72 is_split_into_words=True,
73 return_tensors="pt"
74).to(model.device)
75
76with toch.no_grad():
77 output = model(**inputs)
78 prediction = output.logits.argmax(dim=-1)
79 print(prediction)