Skip to content

Commit

Permalink
Merge pull request #158 from stacklok/bug-109
Browse files Browse the repository at this point in the history
Convert async _close_models in LlamaCppInferenceEngine
  • Loading branch information
lukehinds authored Dec 3, 2024
2 parents b518f0d + fbabe31 commit 12dd7c4
Showing 1 changed file with 10 additions and 10 deletions.
20 changes: 10 additions & 10 deletions src/codegate/inference/inference_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,16 @@ def __init__(self):
self.__models = {}

def __del__(self):
self.__close_models()
self._close_models()

def _close_models(self):
"""
Closes all open models and samplers
"""
for _, model in self.__models.items():
if model._sampler:
model._sampler.close()
model.close()

async def __get_model(self, model_path, embedding=False, n_ctx=512, n_gpu_layers=0):
"""
Expand Down Expand Up @@ -70,12 +79,3 @@ async def embed(self, model_path, content):
"""
model = await self.__get_model(model_path=model_path, embedding=True)
return model.embed(content)

async def __close_models(self):
"""
Closes all open models and samplers
"""
for _, model in self.__models:
if model._sampler:
model._sampler.close()
model.close()

0 comments on commit 12dd7c4

Please sign in to comment.