支持jina-reranker-m0

Author: shell-nlp

gpt_server/model_worker/embedding.py

@@ -49,12 +49,22 @@ def __init__(
                 logger.info(f"device: {self.client.device}")
             self.client.set_processor(model_path)
             self.client.eval()
-        elif "rerank" in self.mode:
+        elif "vl_rerank" == self.mode:
+            self.client = AutoModel.from_pretrained(
+                model_path,
+                torch_dtype="auto",
+                trust_remote_code=True,
+                # attn_implementation="flash_attention_2",
+            )
+
+            self.client.to("cuda")  # or 'cpu' if no GPU is available
+            self.client.eval()
+        elif "rerank" == self.mode:
             self.client = sentence_transformers.CrossEncoder(
                 model_name=model_path, **model_kwargs
             )
             logger.warning("正在使用 rerank 模型...")
-        elif "embedding" in self.mode:
+        elif "embedding" == self.mode:
             self.client = sentence_transformers.SentenceTransformer(
                 model_path, **model_kwargs
             )
@@ -79,6 +89,30 @@ async def get_embeddings(self, params):
             sentence_pairs = [[query, inp] for inp in texts]
             scores = self.client.predict(sentence_pairs)
             embedding = [[float(score)] for score in scores]
+        elif self.mode == "vl_rerank":
+            query = params.get("query", None)
+            token_num = 0
+            sentence_pairs = [[query, inp] for inp in texts]
+            query_type = doc_type = "text"
+            if (
+                query.startswith("http://")
+                or query.startswith("https://")
+                or "data:" in query
+            ):
+                query_type = "image"
+            if (
+                texts[0].startswith("http://")
+                or texts[0].startswith("https://")
+                or "data:" in texts[0]
+            ):
+                doc_type = "image"
+            scores = self.client.compute_score(
+                sentence_pairs,
+                max_length=1024 * 2,
+                query_type=query_type,
+                doc_type=doc_type,
+            )
+            embedding = [[float(score)] for score in scores]
         elif self.mode == "clip_text_model":
             token_num = 0
             if isinstance(texts[0], dict):

gpt_server/model_worker/utils.py

@@ -54,6 +54,10 @@ def get_embedding_mode(model_path: str):
     from infinity_emb.inference.select_model import get_engine_type_from_config
 
     model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    architectures = getattr(model_config, "architectures", [])
+    if "JinaVLForRanking" in architectures:
+        logger.warning("model_type: JinaVLForRanking")
+        return "vl_rerank"
     model_type_text = getattr(
         getattr(model_config, "text_config", {}), "model_type", None
     )
@@ -76,14 +80,13 @@ def get_embedding_mode(model_path: str):
     engine_type_str = str(engine_type)
 
     if "EmbedderEngine" in engine_type_str:
-        mode = "embedding"
+        return "embedding"
     elif "RerankEngine" in engine_type_str:
-        mode = "rerank"
+        return "rerank"
     elif "ImageEmbedEngine" in engine_type_str:
-        mode = model_type or "image"
+        return model_type or "image"
     elif "PredictEngine" in engine_type_str:
-        mode = "classify"
-    return mode
+        return "classify"
 
 
 if __name__ == "__main__":

gpt_server/serving/openai_api_server.py

@@ -134,6 +134,7 @@ async def timing_tasks():
 
     while True:
         try:
+            # ret = await fetch_remote(controller_address + "/refresh_all_workers")
             models = await fetch_remote(
                 controller_address + "/list_models", None, "models"
             )