Skip to content

Commit

Permalink
add model_revison parameter to ModelScopeEmbeddings (langchain-ai#8669)
Browse files Browse the repository at this point in the history

---------

Co-authored-by: Bagatur <[email protected]>
  • Loading branch information
sumandeng and baskaryan authored Aug 3, 2023
1 parent 4e8f11b commit 53e4148
Showing 1 changed file with 8 additions and 5 deletions.
13 changes: 8 additions & 5 deletions libs/langchain/langchain/embeddings/modelscope_hub.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Any, List
from typing import Any, List, Optional

from pydantic import BaseModel, Extra

Expand All @@ -15,27 +15,30 @@ 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"""
super().__init__(**kwargs)
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."""
Expand Down

0 comments on commit 53e4148

Please sign in to comment.