nex/college-bot-v2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Activity Actions

Log download progress to console

Browse source
This commit is contained in:
Nexus 2024-04-14 17:23:05 +01:00
parent 5ac83dd939
commit 086dffd88c
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 10 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/cogs/ollama.py
View file

@ -82,6 +82,13 @@ class OllamaDownloadHandler:
def abort(self): def abort(self):
self._abort.set() self._abort.set()
def __enter__(self):
self._abort.clear()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self._abort.set()
@property @property
def percent(self) -> float: def percent(self) -> float:
return round((self._completed / self._total) * 100, 2) return round((self._completed / self._total) * 100, 2)
@ -946,7 +953,9 @@ class Ollama(commands.Cog):
client = OllamaClient(CONFIG["ollama"][server]["base_url"]) client = OllamaClient(CONFIG["ollama"][server]["base_url"])
if not await client.has_model_named("orca-mini", "3b"): if not await client.has_model_named("orca-mini", "3b"):
await client.download_model("orca-mini", "3b") with client.download_model("orca-mini", "3b") as handler:
async for _ in handler:
self.log.info("Downloading orca-mini:3b on server %r - %.2f%", server, handler.percent)
messages = self.history.get_history(thread) messages = self.history.get_history(thread)
embed = discord.Embed(description=">>> ") embed = discord.Embed(description=">>> ")