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 setuptools>=69
yt-dlp @ https://github.com/yt-dlp/yt-dlp/archive/master.tar.gz yt-dlp @ https://github.com/yt-dlp/yt-dlp/archive/master.tar.gz
# py-cord==2.4.1 # 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 httpx==0.26
psycopg==3.1.16 psycopg==3.1.16
toml==0.10.2 toml==0.10.2
@ -15,3 +15,4 @@ fastapi==0.108.0
uvicorn==0.25.0 uvicorn==0.25.0
psutil==5.9.7 psutil==5.9.7
pydantic==2.5.3 pydantic==2.5.3
humanize==4.9.0

20
src/cogs/ollama.py
View file

@ -7,6 +7,7 @@ import time
import typing import typing
import base64 import base64
import io import io
import humanize
from discord.ui import View, button from discord.ui import View, button
from fnmatch import fnmatch from fnmatch import fnmatch
@ -132,6 +133,7 @@ class Ollama(commands.Cog):
if context not in self.contexts: if context not in self.contexts:
await ctx.respond("Invalid context key.") await ctx.respond("Invalid context key.")
return return
return await ctx.respond("Context is currently disabled.", ephemeral=True)
with open("./assets/ollama-prompt.txt") as file: with open("./assets/ollama-prompt.txt") as file:
system_prompt = file.read() system_prompt = file.read()
await ctx.defer() await ctx.defer()
@ -405,5 +407,23 @@ class Ollama(commands.Cog):
else: else:
await ctx.edit(embed=embed, view=None) 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): def setup(bot):
bot.add_cog(Ollama(bot)) bot.add_cog(Ollama(bot))