add source files
This commit is contained in:
49
barry.py
Normal file
49
barry.py
Normal file
@@ -0,0 +1,49 @@
|
|||||||
|
import re
|
||||||
|
import random
|
||||||
|
|
||||||
|
regex = re.compile("[^a-zA-Z ]")
|
||||||
|
|
||||||
|
class Token:
|
||||||
|
def __init__(self):
|
||||||
|
self.completions: list[str] = []
|
||||||
|
def __repr__(self):
|
||||||
|
return self.completions.__repr__()
|
||||||
|
def add_completion(self, word: str):
|
||||||
|
self.completions.append(word)
|
||||||
|
|
||||||
|
class Model:
|
||||||
|
def __init__(self):
|
||||||
|
self.tokens: dict[str, Token] = {}
|
||||||
|
def predict(self, previous: str) -> str | None:
|
||||||
|
if previous in self.tokens:
|
||||||
|
completion = self.tokens[previous].completions
|
||||||
|
return random.choice(completion)
|
||||||
|
else:
|
||||||
|
return None
|
||||||
|
|
||||||
|
def tokenise(instr: str):
|
||||||
|
return re.split(r" |\n", regex.sub("", instr).lower())
|
||||||
|
|
||||||
|
def train(data: str):
|
||||||
|
model = Model()
|
||||||
|
spldata = tokenise(data)
|
||||||
|
for i in range(0, len(spldata) - 1):
|
||||||
|
word = spldata[i]
|
||||||
|
next_word = spldata[i + 1]
|
||||||
|
if not word in model.tokens:
|
||||||
|
model.tokens[word] = Token()
|
||||||
|
model.tokens[word].add_completion(next_word)
|
||||||
|
return model
|
||||||
|
|
||||||
|
def prompt(prompt: str, model: Model) -> str:
|
||||||
|
words = tokenise(prompt)
|
||||||
|
output: str = ""
|
||||||
|
currentWord = words[-1]
|
||||||
|
for _ in range(1, 200):
|
||||||
|
currentWord = model.predict(currentWord)
|
||||||
|
if currentWord == None:
|
||||||
|
break
|
||||||
|
output += currentWord + " "
|
||||||
|
|
||||||
|
return output
|
||||||
|
|
||||||
28
bot.py
Normal file
28
bot.py
Normal file
@@ -0,0 +1,28 @@
|
|||||||
|
import barry
|
||||||
|
import discord
|
||||||
|
from discord import app_commands
|
||||||
|
import os
|
||||||
|
|
||||||
|
intents = discord.Intents.default()
|
||||||
|
client = discord.Client(intents=intents)
|
||||||
|
tree = app_commands.CommandTree(client)
|
||||||
|
|
||||||
|
with open("training.txt") as t:
|
||||||
|
training = t.read()
|
||||||
|
|
||||||
|
model = barry.train(training)
|
||||||
|
|
||||||
|
@tree.command(name="ask", description="Ask Barry a question")
|
||||||
|
async def ask(interaction: discord.Interaction, question: str):
|
||||||
|
_ = await interaction.response.send_message(barry.prompt(question, model))
|
||||||
|
|
||||||
|
@client.event
|
||||||
|
async def on_ready():
|
||||||
|
await tree.sync()
|
||||||
|
print(f"logged in as {client.user}")
|
||||||
|
key = os.getenv("BARRY_TOKEN", None)
|
||||||
|
if key == None:
|
||||||
|
print("Set BARRY_TOKEN to the Discord token before continuing")
|
||||||
|
exit(1)
|
||||||
|
client.run(key)
|
||||||
|
|
||||||
Reference in New Issue
Block a user