This commit is contained in:
olgavrou 2023-09-04 07:36:47 -04:00
parent ca163f0ee6
commit 67dc1a9dd2
2 changed files with 2 additions and 2 deletions

View File

@ -15,6 +15,7 @@ from langchain.chains.rl_chain.base import (
from langchain.chains.rl_chain.pick_best_chain import ( from langchain.chains.rl_chain.pick_best_chain import (
PickBest, PickBest,
PickBestEvent, PickBestEvent,
PickBestFeatureEmbedder,
PickBestSelected, PickBestSelected,
) )
@ -37,6 +38,7 @@ __all__ = [
"PickBest", "PickBest",
"PickBestEvent", "PickBestEvent",
"PickBestSelected", "PickBestSelected",
"PickBestFeatureEmbedder",
"Embed", "Embed",
"BasedOn", "BasedOn",
"ToSelectFrom", "ToSelectFrom",

View File

@ -60,9 +60,7 @@ class PickBestFeatureEmbedder(base.Embedder[PickBestEvent]):
if model is None: if model is None:
from sentence_transformers import SentenceTransformer from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-mpnet-base-v2") model = SentenceTransformer("all-mpnet-base-v2")
# model = SentenceTransformer("all-MiniLM-L6-v2")
self.model = model self.model = model
self.auto_embed = auto_embed self.auto_embed = auto_embed