From 53e4148a1b3701fdddac2b32a6ce6bb7d46dacfb Mon Sep 17 00:00:00 2001 From: sumandeng Date: Fri, 4 Aug 2023 05:17:48 +0800 Subject: [PATCH] add model_revison parameter to ModelScopeEmbeddings (#8669) --------- Co-authored-by: Bagatur --- .../langchain/embeddings/modelscope_hub.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/libs/langchain/langchain/embeddings/modelscope_hub.py b/libs/langchain/langchain/embeddings/modelscope_hub.py index 6b6ebab42475c..541aeaa358ca7 100644 --- a/libs/langchain/langchain/embeddings/modelscope_hub.py +++ b/libs/langchain/langchain/embeddings/modelscope_hub.py @@ -1,4 +1,4 @@ -from typing import Any, List +from typing import Any, List, Optional from pydantic import BaseModel, Extra @@ -15,12 +15,13 @@ class ModelScopeEmbeddings(BaseModel, Embeddings): from langchain.embeddings import ModelScopeEmbeddings model_id = "damo/nlp_corom_sentence-embedding_english-base" - embed = ModelScopeEmbeddings(model_id=model_id) + embed = ModelScopeEmbeddings(model_id=model_id, model_revision="v1.0.0") """ embed: Any model_id: str = "damo/nlp_corom_sentence-embedding_english-base" """Model name to use.""" + model_revision: Optional[str] = None def __init__(self, **kwargs: Any): """Initialize the modelscope""" @@ -28,14 +29,16 @@ def __init__(self, **kwargs: Any): try: from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks - - self.embed = pipeline(Tasks.sentence_embedding, model=self.model_id) - except ImportError as e: raise ImportError( "Could not import some python packages." "Please install it with `pip install modelscope`." ) from e + self.embed = pipeline( + Tasks.sentence_embedding, + model=self.model_id, + model_revision=self.model_revision, + ) class Config: """Configuration for this pydantic object."""