adds jais2 model support by sarathc-cerebras · Pull Request #42684 · huggingface/transformers

sarathc-cerebras · 2025-12-07T13:43:09Z

What does this PR do?

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Rocketknight1 · 2025-12-08T16:46:01Z

Hi @sarathc-cerebras, thank you for the PR! The main thing missing is a conversion to modular format. You can look at the modular files for other models to see how it works, but it reduces the size of the PR a lot by importing duplicated code from other models.

sarathc-cerebras · 2025-12-09T12:57:31Z

@Rocketknight1 thanks for bringing this up, i have updated it to use the modular format

HuggingFaceDocBuilderDev · 2025-12-09T15:40:24Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Rocketknight1

Yes, this looks good! I made a few comments but they're small.

vasqu

Left some comments, I think we can still simplify a bit and update a few things to be up to date with our current standards. Overall, looking really good already tho

ArthurZucker

LGTM good review @vasqu small nits but let's go!

ArthurZucker · 2025-12-11T08:56:38Z

+        generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+        print(f"Static cache generated text: {generated_text}")
+
+        self.assertGreater(generated_ids.shape[1], input_ids.shape[1])


would be better to have explicit expected outputs here!

vasqu

Please check out the comments from the last review, mostly nits otherwise and let's make the tests more explicit (I've linked an example in one of the review comments)

vasqu

Last comments from my side (I hope), small fixes and finishing touches

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

vasqu · 2025-12-16T13:20:58Z

run-slow: jais2

github-actions · 2025-12-16T13:21:31Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, jais2

github-actions · 2025-12-16T13:22:12Z

This comment contains run-slow, running the specified jobs:

models: ["models/jais2"]
quantizations: []

github-actions · 2025-12-16T13:34:00Z

CI Results

Workflow Run ⚙️

✅ No failing test specific to this PR 🎉 !

github-actions · 2025-12-16T13:55:39Z

View the CircleCI Test Summary for this PR:

https://huggingface.co/spaces/transformers-community/circle-ci-viz?pr=42684&sha=8860cb

sarathc-cerebras · 2025-12-17T10:56:28Z

do i need to do anything else for PR to be merged @vasqu or it will be merged from one of the reviewers?

vasqu · 2025-12-17T11:08:14Z

@sarathc-cerebras merged now, sorry about that. The CI blocked but it isn't related to your PR 🤗

* adds jais2 model support * updates tests * addresses review comment * review comments addressed * addresses test review comments * fixes date * format issue fix * Update src/transformers/models/jais2/__init__.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update tests/models/jais2/test_modeling_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * Update src/transformers/models/jais2/modular_jais2.py Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * fixes tests as per review comment * updates layernorm setup * Apply suggestions from code review Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> * addressed review comments and updated tests as recomended * fixup tests --------- Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com> Co-authored-by: vasqu <antonprogamer@gmail.com>

sarathc-cerebras force-pushed the add-jais2-model branch from 377e2b8 to ab785fc Compare December 9, 2025 12:01

sarathc-cerebras force-pushed the add-jais2-model branch 4 times, most recently from 2ae7204 to 672e38a Compare December 9, 2025 14:13

sarathc-cerebras force-pushed the add-jais2-model branch from 9e0839b to 7dfa45e Compare December 9, 2025 16:10

Rocketknight1 approved these changes Dec 9, 2025

View reviewed changes

Comment thread src/transformers/models/jais2/modular_jais2.py Outdated

Comment thread docs/source/en/model_doc/jais2.md Outdated

Comment thread tests/models/jais2/test_modeling_jais2.py Outdated

Comment thread docs/source/en/model_doc/jais2.md Outdated

sarathc-cerebras force-pushed the add-jais2-model branch from a363e45 to e363470 Compare December 9, 2025 16:58

vasqu reviewed Dec 10, 2025

View reviewed changes

sarathc-cerebras force-pushed the add-jais2-model branch from 2f9713c to 5090c18 Compare December 10, 2025 14:25

ArthurZucker approved these changes Dec 11, 2025

View reviewed changes

sarathc-cerebras force-pushed the add-jais2-model branch from efed368 to f4a67f3 Compare December 11, 2025 12:01

vasqu reviewed Dec 11, 2025

View reviewed changes

sarathc-cerebras force-pushed the add-jais2-model branch from 10a5980 to bf97684 Compare December 12, 2025 19:13

vasqu approved these changes Dec 15, 2025

View reviewed changes

sarathc-cerebras and others added 12 commits December 16, 2025 12:01

adds jais2 model support

c65baec

updates tests

ce372f4

addresses review comment

e6ad269

review comments addressed

5ae5c35

addresses test review comments

a2df8cd

fixes date

18adadc

format issue fix

afc73f5

Update src/transformers/models/jais2/__init__.py

63f1262

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update src/transformers/models/jais2/modular_jais2.py

aa63e5c

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update tests/models/jais2/test_modeling_jais2.py

78924e1

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update src/transformers/models/jais2/modular_jais2.py

ce62466

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update src/transformers/models/jais2/modular_jais2.py

a553d2b

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

sarathc-cerebras and others added 5 commits December 16, 2025 12:01

Update src/transformers/models/jais2/modular_jais2.py

8223e31

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update src/transformers/models/jais2/modular_jais2.py

36fb5f5

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Update src/transformers/models/jais2/modular_jais2.py

8348b3f

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

fixes tests as per review comment

5c60f4a

updates layernorm setup

506d4fd

sarathc-cerebras force-pushed the add-jais2-model branch from 9398ddb to 506d4fd Compare December 16, 2025 08:01

sarathc-cerebras and others added 3 commits December 16, 2025 13:32

Apply suggestions from code review

6ec0bc6

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

addressed review comments and updated tests as recomended

a29243e

fixup tests

c25937e

Merge branch 'main' into add-jais2-model

8860cb4

vasqu enabled auto-merge (squash) December 16, 2025 13:48

sarathc-cerebras mentioned this pull request Dec 16, 2025

[Model] adds jais 2 support vllm-project/vllm#30188

Merged

5 tasks

ArthurZucker disabled auto-merge December 17, 2025 11:06

ArthurZucker merged commit 0dbf808 into huggingface:main Dec 17, 2025
22 of 25 checks passed

Conversation

sarathc-cerebras commented Dec 7, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

Rocketknight1 commented Dec 8, 2025

Uh oh!

sarathc-cerebras commented Dec 9, 2025

Uh oh!

HuggingFaceDocBuilderDev commented Dec 9, 2025

Uh oh!

Rocketknight1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ArthurZucker left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

ArthurZucker Dec 11, 2025

Choose a reason for hiding this comment

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vasqu commented Dec 16, 2025

Uh oh!

github-actions Bot commented Dec 16, 2025

Uh oh!

github-actions Bot commented Dec 16, 2025

Uh oh!

github-actions Bot commented Dec 16, 2025

CI Results

Uh oh!

github-actions Bot commented Dec 16, 2025

Uh oh!

sarathc-cerebras commented Dec 17, 2025

Uh oh!

Uh oh!

vasqu commented Dec 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

sarathc-cerebras commented Dec 7, 2025 •

edited

Loading