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

Add timing info to ollama

Browse source
This commit is contained in:
Nexus 2024-01-12 09:53:57 +00:00
parent 0193ab25ad
commit 4cfb598e4d
2 changed files with 22 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
requirements.txt
View file

@ -2,7 +2,7 @@ wheel>=0.42
setuptools>=69
yt-dlp @ https://github.com/yt-dlp/yt-dlp/archive/master.tar.gz
# py-cord==2.4.1
py-cord @ git+https://github.com/Pycord-Development/pycord.git
py-cord @ git+https://github.com/Pycord-Development/pycord.git@0fe9265665680ba7169128849ab9fdb48b8a5954
httpx==0.26
psycopg==3.1.16
toml==0.10.2
@ -15,3 +15,4 @@ fastapi==0.108.0
uvicorn==0.25.0
psutil==5.9.7
pydantic==2.5.3
humanize==4.9.0

20
src/cogs/ollama.py
View file

@ -7,6 +7,7 @@ import time
import typing
import base64
import io
import humanize
from discord.ui import View, button
from fnmatch import fnmatch
@ -132,6 +133,7 @@ class Ollama(commands.Cog):
if context not in self.contexts:
await ctx.respond("Invalid context key.")
return
return await ctx.respond("Context is currently disabled.", ephemeral=True)
with open("./assets/ollama-prompt.txt") as file:
system_prompt = file.read()
await ctx.defer()
@ -405,5 +407,23 @@ class Ollama(commands.Cog):
else:
await ctx.edit(embed=embed, view=None)
if line.get("done"):
total_duration = humanize.naturaldelta(line["total_duration"] / 1e9)
load_duration = humanize.naturaldelta(line["load_duration"] / 1e9)
prompt_eval_duration = humanize.naturaldelta(line["prompt_eval_duration"] / 1e9)
eval_duration = humanize.naturaldelta(line["eval_duration"] / 1e9)
embed = discord.Embed(
title="Timings",
description=f"Total: {total_duration}\nLoad: {load_duration}\n"
f"Prompt Eval: {prompt_eval_duration}\nEval: {eval_duration}\n"
f"Prompt Tokens: {line['prompt_eval_count']:,}\n"
f"Response Tokens: {line['eval_count']:,}",
color=discord.Color.blurple(),
timestamp=discord.utils.utcnow()
)
return await ctx.respond(embed=embed, ephemeral=True)
def setup(bot):
bot.add_cog(Ollama(bot))