Skip to content

Commit 546e8de

Browse files
committed
Extract model loading into separate function
In order to properly include the model in the image, the download needs to be individually callable without the whole application.
1 parent f319825 commit 546e8de

2 files changed

Lines changed: 10 additions & 2 deletions

File tree

src/endpoint/generation/generate_endpoint.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
from litellm.types.completion import ChatCompletionUserMessageParam
77
from pydantic import ValidationError
88
from qdrant_client import QdrantClient
9-
from sentence_transformers import SentenceTransformer
9+
10+
from src.model import load_vector_model
1011

1112
from src.mapper.data_type_mapper import map_to_data_type_schema
1213
from src.mapper.flow_mapper import map_to_grpc_flow, map_to_flow_schema
@@ -25,7 +26,7 @@ class GenerateService(pb2_grpc.GenerateServiceServicer):
2526

2627
def __init__(self):
2728
self.memory_client = QdrantClient(":memory:")
28-
self.vector_model = SentenceTransformer('all-MiniLM-L6-v2')
29+
self.vector_model = load_vector_model()
2930
self.function_store = FunctionStore(self.memory_client, self.vector_model)
3031
self.flow_type_store = FlowTypeStore(self.memory_client, self.vector_model)
3132
self.model_store = ModelStore()

src/model.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
from sentence_transformers import SentenceTransformer
2+
3+
MODEL_NAME = 'all-MiniLM-L6-v2'
4+
5+
6+
def load_vector_model() -> SentenceTransformer:
7+
return SentenceTransformer(MODEL_NAME)

0 commit comments

Comments
 (0)