add model_revison parameter to ModelScopeEmbeddings (#8669)

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
This commit is contained in:
sumandeng 2023-08-04 05:17:48 +08:00 committed by GitHub
parent 4e8f11b36a
commit 53e4148a1b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -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 @@ class ModelScopeEmbeddings(BaseModel, Embeddings):
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."""