fix - some qwen3_reranker have no lm_weight

fix - support qwen3 rerank after refactor embedding endpoint

Author: jianglan89

rtp_llm/models/downstream_modules/reranker/qwen3_reranker.py

@@ -107,11 +107,15 @@ def __init__(
         super().__init__(config)
         self.token_false_id = token_false_id
         self.token_true_id = token_true_id
-
+        self.tie_word_embeddings = config.tie_word_embeddings
+        self.lm_head_weight_name = (
+            "model.embed_tokens.weight"
+            if self.tie_word_embeddings
+            else "lm_head.weight"
+        )
+        
     def custom_weight_info(self) -> List[CustomAtomicWeight]:
-        w_list = [
-            "lm_head.weight",
-        ]
+        w_list = [self.lm_head_weight_name]
         weights = []
         for k in w_list:
             weights.append(
@@ -121,7 +125,7 @@ def custom_weight_info(self) -> List[CustomAtomicWeight]:
 
     def init(self, tensor_map: Dict[str, torch.Tensor]):
         data_type = to_torch_dtype(self.config_.data_type)
-        linear_weight = tensor_map["lm_head.weight"]
+        linear_weight = tensor_map[self.lm_head_weight_name]
         self.linear = torch.nn.Linear(linear_weight.shape[1], linear_weight.shape[0])
         self.linear.weight.data = linear_weight
         self.linear = self.linear.to(data_type).eval().to(self.device)

rtp_llm/models/downstream_modules/utils.py

@@ -2,7 +2,6 @@
 from typing import Optional
 
 from rtp_llm.config.model_config import ModelConfig
-from rtp_llm.ops import TaskType
 from rtp_llm.frontend.tokenizer_factory.tokenizers import BaseTokenizer
 from rtp_llm.models.downstream_modules import (
     ALLEmbeddingModule,
@@ -13,27 +12,32 @@
     RerankerModule,
     SparseEmbeddingModule,
 )
+from rtp_llm.models.downstream_modules.reranker.qwen3_reranker import (
+    Qwen3RerankerModule,
+)
+from rtp_llm.ops import TaskType
 
 
 def create_custom_module(
     config: ModelConfig,
     tokenizer: Optional[BaseTokenizer],
-):  
+):
     # try import internal module
     try:
         from internal_source.rtp_llm.models.downstream_modules.utils import (
             create_custom_module,
         )
+
         internal_module = create_custom_module(config, tokenizer)
         if internal_module is not None:
             return internal_module
     except ImportError:
         logging.exception("internal module not found, using external module")
 
-
     task_type = config.task_type
     if task_type == TaskType.LANGUAGE_MODEL:
         return None
+    model_type = config.model_type
     assert tokenizer is not None, "tokenizer should not be None"
     if task_type == TaskType.DENSE_EMBEDDING:
         return DenseEmbeddingModule(config, tokenizer)
@@ -47,6 +51,8 @@ def create_custom_module(
         return ClassifierModule(config, tokenizer)
     elif task_type == TaskType.BGE_M3:
         return BgeM3EmbeddingModule(config, tokenizer)
+    elif model_type == "qwen_3":
+        return Qwen3RerankerModule(config, tokenizer)
     elif task_type == TaskType.RERANKER:
         return RerankerModule(config, tokenizer)
     raise Exception(f"unknown task_type: {task_type}")