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

Update ollama progress bar
All checks were successful
Build and Publish Jimmy.2 / build_and_publish (push) Successful in 5s
Details

Browse source
This commit is contained in:
Nexus 2024-05-31 18:24:15 +01:00
parent 70afac54ca
commit 09f8a59d6b
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 8 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
src/cogs/ollama.py
View file

@ -654,10 +654,12 @@ class Ollama(commands.Cog):
if resp.status == 404: if resp.status == 404:
self.log.debug("Beginning download of %r", model) self.log.debug("Beginning download of %r", model)
def progress_bar(_v: float, action: str = None): def progress_bar(_v: float, action: str = None, _mbps: float = None):
bar = "\N{large green square}" * round(_v / 10) bar = "\N{large green square}" * round(_v / 10)
bar += "\N{white large square}" * (10 - len(bar)) bar += "\N{white large square}" * (10 - len(bar))
bar += f" {_v:.2f}%" bar += f" {_v:.2f}%"
if _mbps:
bar += f" ({_mbps:.2f} MB/s)"
if action: if action:
return f"{action} {bar}" return f"{action} {bar}"
return bar return bar
@ -685,6 +687,7 @@ class Ollama(commands.Cog):
embed.set_footer(text="Unable to continue.") embed.set_footer(text="Unable to continue.")
return await ctx.edit(embed=embed) return await ctx.edit(embed=embed)
view = OllamaView(ctx) view = OllamaView(ctx)
last_downloaded = 0
async for line in ollama_stream(response.content): async for line in ollama_stream(response.content):
if view.cancel.is_set(): if view.cancel.is_set():
embed = discord.Embed( embed = discord.Embed(
@ -695,11 +698,14 @@ class Ollama(commands.Cog):
return await ctx.edit(embed=embed, view=None) return await ctx.edit(embed=embed, view=None)
if time.time() >= (last_update + 5.1): if time.time() >= (last_update + 5.1):
if line.get("total") is not None and line.get("completed") is not None: if line.get("total") is not None and line.get("completed") is not None:
new_bytes = line["completed"] - last_downloaded
mbps = new_bytes / 1024 / 1024 / 5
percent = (line["completed"] / line["total"]) * 100 percent = (line["completed"] / line["total"]) * 100
else: else:
percent = 50.0 percent = 50.0
mbps = 0.0
embed.fields[0].value = progress_bar(percent, line["status"]) embed.fields[0].value = progress_bar(percent, line["status"], mbps)
await ctx.edit(embed=embed, view=view) await ctx.edit(embed=embed, view=view)
last_update = time.time() last_update = time.time()
else: else: