Add attention and final logit softcapping.

This commit is contained in:
Andrei Betlen 2024-06-28 15:42:19 -04:00
parent 8748d8ac6f
commit 4d3f17b4ac
4 changed files with 38 additions and 2 deletions

View file

@ -50,6 +50,8 @@ class Keys:
POOLING_TYPE = "{arch}.pooling_type"
LOGIT_SCALE = "{arch}.logit_scale"
DECODER_START_TOKEN_ID = "{arch}.decoder_start_token_id"
ATTN_LOGIT_SOFTCAPPING = "{arch}.attn_logit_softcapping"
FINAL_LOGIT_SOFTCAPPING = "{arch}.final_logit_softcapping"
class Attention:
HEAD_COUNT = "{arch}.attention.head_count"