jalr/tabbyAPI-ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

[fix] Bring draft progress and model progress in sync with model loader (#125)

Browse source 
* Bring draft progress and model progress in sync with model loader

* Fix formatting
This commit is contained in:
turboderp 2024-06-03 19:41:02 +02:00 • committed by GitHub
parent a011c17488
commit 0eb8fa5d1e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

5
common/model.py
View file

@ -62,14 +62,15 @@ async def load_model_gen(model_path: pathlib.Path, **kwargs):
)
else:
progress.advance(loading_task)
yield module, modules, model_type
if module == modules:
# Switch to model progress if the draft model is loaded
if model_type == "draft":
model_type = "model"
else:
progress.stop()
yield module, modules, model_type
finally:
progress.stop()