Skip to content

[Feature] Add support for Qwen3 Reranker with Sequence Classifier head #698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

danielchalef
Copy link

@danielchalef danielchalef commented Aug 13, 2025

I totally missed #695 and ended up coding this in a hurry today.

@OlivierDehaene or @Narsil: Happy to work with @sigridjineth if helpful

Summary

Adds support for Qwen3 reranker models (e.g., tomaarsen/Qwen3-Reranker-0.6B-seq-cls) that have been converted from causal LM to sequence classification architecture. These models require specific chat templates for optimal performance. Functionality was added to support custom instructions.

How Qwen3 Reranking (with Seq Classification conversion) Works

Conversion adds a classification head (score.weight) on top of base model

  • This PR enables TEI to recognize these models as classifier/reranker
  • Works with /rerank endpoint
  • Process:
    a. Takes query-document pair formatted as conversation
    b. Uses LastToken pooling to extract representations
    c. Passes through classification head (single linear layer, no bias)
    d. Returns logit score (sigmoid applied for probability)

Key Changes

  1. Template Support
  • New templates module implementing Qwen3 chat format:
    <|im_start|>system Judge whether the Document meets the requirements... <|im_start|>user <Instruct>: {instruction} <Query>: {query} <Document>: {document}<|im_end|> <|im_start|>assistant <think>
  • Added instruction and use_template fields to RerankRequest (HTTP & gRPC)
  • Auto-detects Qwen3 seq-cls models for template application
  • Default instruction: "Select only the Documents that are semantically similar to the Query."
  1. Model Loading
  • Flexible tensor loading for score.weight pattern (no bias tensor)
  • Uses LastToken pooling for classification models
  • Available via /rerank endpoint for Classifier models

Testing

  • Added test_qwen3_reranker() test in backends/candle/tests/test_qwen3.rs
  • Test verifies model loading and prediction functionality

Notes

The model is highly sensitive to the instruction and template format. Without the template, results can be somewhat random. Tuning the Instruction can result in significantly improved discrimination.

danielchalef and others added 2 commits August 12, 2025 16:56
- Added Qwen3ClassificationHead with flexible tensor loading that handles:
  - score.weight at top level (for converted Qwen3 rerankers)
  - classifier.weight/bias patterns for standard models
- Updated Qwen3Model and FlashQwen3Model to support classification
- Added predict() method implementations for both model variants
- Extended Qwen3Config with id2label field for classification
- Added test case for Qwen3 reranker models with snapshot

The implementation supports Qwen3 models converted to sequence classifiers
for reranking tasks (e.g., tomaarsen/Qwen3-Reranker-0.6B-seq-cls).
The classification head gracefully handles different tensor naming
conventions from various conversion approaches.

Tested with both embedding and reranking Qwen3 models.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <[email protected]>
@sigridjineth
Copy link

sigridjineth commented Aug 13, 2025

Hello, Thanks for your contribution. #695 was my first work to TEI so I want to merge my contribution to your PR if possible.

Do you only support sep-cls tuned model or do support CasualModel as well? Please note that in #695, I focused on supporting Qwen3CasualModel not a classification tuned model so that I added a new type of ModelType (called ListWiseReranker)

@danielchalef
Copy link
Author

danielchalef commented Aug 13, 2025

@sigridjineth this PR only supports seq cls models. It was the lowest-effort way of implementing support for Qwen3-Reranker and also far more efficient. The classification head is less computationally costly than the causal LM head by orders of magnitude.

I'm happy for you to take this PR and merge it with yours, if you'd like. This way we can support both forms of models.

Copy link

@sigridjineth sigridjineth left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@danielchalef hello, when requesting like below, I got the issue.

curl --location 'http://localhost:8080/rerank' \
--header 'Content-Type: application/json' \
--header 'Cookie: locale=en' \
--data '{
    "query": "털사 대학교에서 2003년부터 2006년까지 감독을 맡았던 사람이 누구야?",
    "texts": [
        "존 맥널티(1968년 5월 29일 출생)는 미국 프로 미식축구 리그(NFL) 로스앤젤레스 차저스의 타이트 엔드 코치인 미식축구 코치입니다. 그는 1990년 펜실베이니아 대학교를 졸업했습니다. 2012년까지 애리조나 카디널스에서 코치로 활동했으며, 2009년 초 와이드 리시버 코치로 시작하여 2012년 켄 위젠헌트 감독이 경질될 때까지 쿼터백 코치로 활동했습니다. 이후 2013년에는 그렉 샤이아노 감독이 경질될 때까지 탬파베이 버커니어스의 쿼터백 코치로 활동했습니다. 그리고 2014년에는 테네시 타이탄스의 쿼터백 코치로 합류하여 마이크 뮬라키 감독이 켄 위젠헌트 감독을 경질한 2015년까지 그곳에서 일했습니다. 현재 그는 로스앤젤레스 차저스의 타이트 엔드 코치로 일하고 있습니다.",
        "브라이언 쇼튼하이머(1973년 10월 16일 출생)는 미국 프로 미식축구 리그(NFL) 인디애나폴리스 콜츠의 쿼터백 코치인 미국 미식축구 코치입니다. 그는 이전에 NFL의 워싱턴 레드스킨스와 샌디에이고 차저스에서 쿼터백 코치로 활약했으며, NFL의 뉴욕 제츠, 세인트루이스 램스, 조지아 대학교의 조지아 불독스 미식축구 팀의 공격 코디네이터를 역임했습니다. 그의 아버지인 마티 쇼튼하이머는 전 캔자스시티 치프스 감독이었고 그의 삼촌인 커트 쇼튼하이머도 치프스의 수비 백 코치였습니다.",
        "랜디 샌더스(1965년 9월 22일 출생)는 현재 플로리다 주립대학교의 쿼터백 코치 겸 공동 공격 코디네이터로 활동하고 있는 미국의 미식축구 코치입니다. 그는 최초의 우승팀인 1998년 테네시 대학교의 쿼터백 코치 겸 공격 코디네이터와 마지막 우승팀인 2013년 플로리다 주립대의 쿼터백 코치로서 활약한 유일한 인물입니다.",
        "메이저 리 애플화이트(1978년 7월 26일 출생)는 미국의 미식축구 코치이자 전 선수입니다. 그는 현재 휴스턴 대학교의 감독이며, 이전에는 공격 코디네이터와 쿼터백 코치로 활약했습니다. 2013년에는 텍사스 대학교에서 공동 공격 코디네이터 겸 쿼터백 코치로 활동했습니다. 텍사스로 오기 전에는 2006년 토드 그레이엄 감독 밑에서 라이스 대학교의 공격 코디네이터로, 2007년에는 닉 새빈 감독 밑에서 앨라배마 대학교의 공격 코디네이터로 일했습니다. 당시 그는 디비전 I-A 학교 중 최연소 공격 코디네이터였습니다.",
        "모리스 왓츠(1936년생)는 은퇴한 미국 미식축구 코치이자 전 선수입니다. 2016년 은퇴할 때까지 센트럴 미시간 대학교에서 공격 코디네이터와 쿼터백 코치로 활약했습니다. 그는 마이애미 대학교에서도 공격 코디네이터로 재직했으며, 2009년 쿼터백 코치로 레드호크스 코칭 스태프에 합류했습니다. 모리스는 2002 시즌 바비 윌리엄스 감독이 경질된 후 마지막 세 경기 동안 미시간 주립대학교의 임시 감독을 맡아 1승 2패의 성적을 기록했습니다.",
        "스티븐 존 크래고퍼(1965년 4월 28일 출생)는 미식축구 코치이자 전 선수입니다. 그는 이전에 타이거즈 풋볼 팀의 쿼터백 코치로 일한 후 현재 루이지애나 주립대학교 미식축구 프로그램의 행정 보좌관으로 일하고 있습니다. 그는 2007년부터 2009년까지 루이빌에서, 2003년부터 2006년까지 툴사에서 헤드 코치로 역임했습니다.",
        "11월 1일생(1974년)인 스콧 로펠러는 미식축구 코치이자 전 선수입니다. 그는 현재 보스턴 칼리지에서 공격 코디네이터 겸 쿼터백 코치로 활동하고 있습니다. 이전에는 버지니아 공과대학에서 공격 코디네이터와 쿼터백 코치로 활약했으며, 진 치직 감독 밑에서 오번 대학교에서도 같은 역할을 맡았습니다. 오번 대학교에 합류하기 전에는 템플 대학교에서 공격 코디네이터로 재직했습니다. 그는 주로 빅 텐과 서던 컨퍼런스에서 10년 이상 쿼터백 코치로 경력을 쌓았습니다.",
        "밥 샌더스는 현재 내셔널 풋볼 리그(NFL)의 클리블랜드 브라운스에서 공격 보조 코치로 활동하고 있습니다. 그는 전 샌디에이고 차저스 감독인 알 샌더스의 아들입니다. 샌더스는 딕 버밀 감독이 이끄는 캔자스시티 치프스에서 NFL 최고 공격진의 보조 코치로 활약했으며, 명예의 전당에 오른 조 깁 감독 밑에서 워싱턴 레드스킨스와 세인트루이스 램스에서 뛰었습니다. 그는 미국 챔피언십 버지니아 디스트로이어스 프랜차이즈에서 마티 쇼트니히머 감독과 함께 두 시즌 동안 와이드 리시버 코치로 활약했습니다. 대학에서는 캔자스 웨슬리언 대학교와 오버린 칼리지에서 공격 코디네이터 겸 쿼터백 코치로, 워싱턴 대학교 세인트루이스에서 쿼터백 코치로 재직했습니다. 샌더스는 미국 최고의 대학 선수들이 출전하는 NFLPA 올스타 게임의 리시버 코치로 세 차례 선정된 바 있습니다. 샌더스는 남부 감리교 대학교에서 대학 미식축구 선수로 활약했습니다.",
        "마이크 오케인(1954년 7월 20일 출생)은 미국의 미식축구 코치이자 전 선수입니다. 현재 이스트 테네시 주립대학교 미식축구팀의 공격 코디네이터와 쿼터백 코치를 맡고 있습니다. 그는 제임스 매디슨의 전 공격 코디네이터였으며 2006 시즌부터 2012 시즌 종료까지 버지니아 공과대학 호키스 미식축구팀의 쿼터백 코치를 역임했습니다. 오케인은 1993년부터 1999년까지 노스캐롤라이나주립대학교 미식축구팀 감독을 맡아 41승 40패의 성적을 기록했습니다. 그는 1974년부터 1976년까지 클렘슨 대학교에서 쿼터백과 펀터로 미식축구 선수로 활약했습니다.",
        "2005년 툴사 골든 허리케인 축구팀은 2005년 NCAA 디비전 I-A 축구 시즌에서 툴사 대학교를 대표했습니다. 팀의 감독은 스티브 크래고였습니다. 이들은 오클라호마 툴사에 있는 켈리 경기장에서 홈 경기를 치렀고 콘퍼런스 USA 서부 지구에서 경쟁했습니다."
    ],
    "truncate": false
}'

2025-08-13T17:38:17.898462Z ERROR rerank:predict{truncate=false truncation_direction=Right raw_scores=false}: text_embeddings_core::infer: core/src/infer.rs:450: unexpected rank, expected: 2, got: 3 ([8, 8, 1])


{
    "error": "Backend error: unexpected rank, expected: 2, got: 3 ([8, 8, 1])",
    "error_type": "Backend"
}

doesn't we need to select the first token from each sequence in the batch for CLS pooling, in impl Qwen3Model ?

@danielchalef
Copy link
Author

danielchalef commented Aug 13, 2025

@sigridjineth Oh wow. I somehow neglected to commit the entirety of my working branch. :-/ I'm so sorry. Can you try again? You'll note this now in qwen3.rs:

let classification_head = Some(Qwen3ClassificationHead::load(vb.clone(), config)?);
                (Pool::LastToken, classification_head) 

That's odd. Using tomaarsen/Qwen3-Reranker-0.6B-seq-cls, I get the following from a container built on this branch.

Which model are you using?

curl --location 'http://XXXX.getzep.com/rerank' \
--header 'Content-Type: application/json' \
--header 'Cookie: locale=en' \
--data '{
    "query": "털사 대학교에서 2003년부터 2006년까지 감독을 맡았던 사람이 누구야?",
    "texts": [
        "존 맥널티(1968년 5월 29일 출생)는 미국 프로 미식축구 리그(NFL) 로스앤젤레스 차저스의 타이트 엔드 코치인 미식축구 코치입니다. 그는 1990년 펜실베이니아 대학교를 졸업했습니다. 2012년까지 애리조나 카디널스에서 코치로 활동했으며, 2009년 초 와이드 리시버 코치로 시작하여 2012년 켄 위젠헌트 감독이 경질될 때까지 쿼터백 코치로 활동했습니다. 이후 2013년에는 그렉 샤이아노 감독이 경질될 때까지 탬파베이 버커니어스의 쿼터백 코치로 활동했습니다. 그리고 2014년에는 테네시 타이탄스의 쿼터백 코치로 합류하여 마이크 뮬라키 감독이 켄 위젠헌트 감독을 경질한 2015년까지 그곳에서 일했습니다. 현재 그는 로스앤젤레스 차저스의 타이 트 엔드 코치로 일하고 있습니다.",
        "브라이언 쇼튼하이머(1973년 10월 16일 출생)는 미국 프로 미식축구 리그(NFL) 인디애나폴리스 콜츠 의 쿼터백 코치인 미국 미식축구 코치입니다. 그는 이전에 NFL의 워싱턴 레드스킨스와 샌디에이고 차저스에서 쿼터백 코치로 활약했으며, NFL의 뉴욕 제츠, 세인트루이스 램스, 조지아 대학교의 조지아 불독스 미식축구 팀의 공격 코디네이터를 역임했습니다. 그의 아버지인 마티 쇼튼하이머는 전 캔자스시티 치프스 감독이었고 그의 삼촌인 커트 쇼튼하이머도 치프스의 수비 백 코치였습니다.",
        "랜디 샌더스(1965년 9월 22일 출생)는 현재 플로리다 주립대학교의 쿼터백 코치 겸 공동 공격 코디네이터로 활동하고 있는 미국의 미식축구 코치입니다. 그는 최초의 우승팀인 1998년 테네시 대학교의 쿼터백 코 치 겸 공격 코디네이터와 마지막 우승팀인 2013년 플로리다 주립대의 쿼터백 코치로서 활약한 유일한 인물입니다.",
        "메이저 리 애플화이트(1978년 7월 26일 출생)는 미국의 미식축구 코치이자 전 선수입니다. 그는 현재 휴스턴 대학교의 감독이며, 이전에는 공격 코디네이터와 쿼터백 코치로 활약했습니다. 2013년에는 텍사스 대 학교에서 공동 공격 코디네이터 겸 쿼터백 코치로 활동했습니다. 텍사스로 오기 전에는 2006년 토드 그레이엄 감독 밑에서 라이스 대학교의 공격 코디네이터로, 2007년에는 닉 새빈 감독 밑에서 앨라배마 대학교의 공격 코디네이터로 일했습니다. 당시 그는 디비전 I-A 학교 중 최연소 공격 코디네이터였습니다.",
        "모리스 왓츠(1936년생)는 은퇴한 미국 미식축구 코치이자 전 선수입니다. 2016년 은퇴할 때까지 센트럴 미시간 대학교에서 공격 코디네이터와 쿼터백 코치로 활약했습니다. 그는 마이애미 대학교에서도 공격 코디네이터로 재직했으며, 2009년 쿼터백 코치로 레드호크스 코칭 스태프에 합류했습니다. 모리스는 2002 시즌 바 비 윌리엄스 감독이 경질된 후 마지막 세 경기 동안 미시간 주립대학교의 임시 감독을 맡아 1승 2패의 성적을 기록했습니다.",
        "스티븐 존 크래고퍼(1965년 4월 28일 출생)는 미식축구 코치이자 전 선수입니다. 그는 이전에 타이거즈 풋볼 팀의 쿼터백 코치로 일한 후 현재 루이지애나 주립대학교 미식축구 프로그램의 행정 보좌관으로 일하 고 있습니다. 그는 2007년부터 2009년까지 루이빌에서, 2003년부터 2006년까지 툴사에서 헤드 코치로 역임했습니다.",
        "11월 1일생(1974년)인 스콧 로펠러는 미식축구 코치이자 전 선수입니다. 그는 현재 보스턴 칼리지에 서 공격 코디네이터 겸 쿼터백 코치로 활동하고 있습니다. 이전에는 버지니아 공과대학에서 공격 코디네이터와 쿼터백 코치로 활약했으며, 진 치직 감독 밑에서 오번 대학교에서도 같은 역할을 맡았습니다. 오번 대학교에 합류하기 전에는 템플 대학교에서 공격 코디네이터로 재직했습니다. 그는 주로 빅 텐과 서던 컨퍼런스에서 10 년 이상 쿼터백 코치로 경력을 쌓았습니다.",
        "밥 샌더스는 현재 내셔널 풋볼 리그(NFL)의 클리블랜드 브라운스에서 공격 보조 코치로 활동하고 있 습니다. 그는 전 샌디에이고 차저스 감독인 알 샌더스의 아들입니다. 샌더스는 딕 버밀 감독이 이끄는 캔자스 시티 치프스에서 NFL 최고 공격진의 보조 코치로 활약했으며, 명예의 전당에 오른 조 깁 감독 밑에서 워싱턴  레드스킨스와 세인트루이스 램스에서 뛰었습니다. 그는 미국 챔피언십 버지니아 디스트로이어스 프랜차이즈에 서 마티 쇼트니히머 감독과 함께 두 시즌 동안 와이드 리시버 코치로 활약했습니다. 대학에서는 캔자스 웨슬리언 대학교와 오버린 칼리지에서 공격 코디네이터 겸 쿼터백 코치로, 워싱턴 대학교 세인트루이스에서 쿼터백  코치로 재직했습니다. 샌더스는 미국 최고의 대학 선수들이 출전하는 NFLPA 올스타 게임의 리시버 코치로 세  차례 선정된 바 있습니다. 샌더스는 남부 감리교 대학교에서 대학 미식축구 선수로 활약했습니다.",
        "마이크 오케인(1954년 7월 20일 출생)은 미국의 미식축구 코치이자 전 선수입니다. 현재 이스트 테네시 주립대학교 미식축구팀의 공격 코디네이터와 쿼터백 코치를 맡고 있습니다. 그는 제임스 매디슨의 전 공격 코디네이터였으며 2006 시즌부터 2012 시즌 종료까지 버지니아 공과대학 호키스 미식축구팀의 쿼터백 코치를  역임했습니다. 오케인은 1993년부터 1999년까지 노스캐롤라이나주립대학교 미식축구팀 감독을 맡아 41승 40패 의 성적을 기록했습니다. 그는 1974년부터 1976년까지 클렘슨 대학교에서 쿼터백과 펀터로 미식축구 선수로 활약했습니다.",
        "2005년 툴사 골든 허리케인 축구팀은 2005년 NCAA 디비전 I-A 축구 시즌에서 툴사 대학교를 대표했습니다. 팀의 감독은 스티브 크래고였습니다. 이들은 오클라호마 툴사에 있는 켈리 경기장에서 홈 경기를 치렀고 콘퍼런스 USA 서부 지구에서 경쟁했습니다."
    ],
    "truncate": false
}'

[{"index":5,"score":0.9554855},{"index":8,"score":0.0001313518},{"index":3,"score":0.00011147904},{"index":9,"score":0.000058291265},{"index":6,"score":0.000036478537},{"index":7,"score":0.0000059551394},{"index":4,"score":0.000005338156},{"index":0,"score":0.000005214498},{"index":2,"score":0.00000283508},{"index":1,"score":0.0000012878952}]%                                                                   

- Added optional fields `instruction` and `use_template` to `RerankRequest` for custom instructions and template usage.
- Updated rerank logic to apply templates conditionally based on model type and user input.
- Introduced template formatting for Qwen3 models to improve reranking results.

This update allows for more flexible and context-aware reranking capabilities, particularly for Qwen3 models.
@sigridjineth
Copy link

@danielchalef have you tried in cuda devices? it got the same issue that gives the same score across the inputs.

@sigridjineth
Copy link

@danielchalef made fix on the flashqwen3.rs danielchalef#1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants