104 lines
2.8 KiB
Python
104 lines
2.8 KiB
Python
import random
|
|
import os
|
|
|
|
from rich import traceback
|
|
|
|
import model
|
|
import utils
|
|
from client import DramaClient
|
|
|
|
from config import config
|
|
|
|
traceback.install()
|
|
|
|
|
|
class Bot:
|
|
def __init__(self):
|
|
print("Loading model...")
|
|
self.model = model.Model()
|
|
self.client = DramaClient()
|
|
print("Ready.")
|
|
|
|
def post_random_reply(self):
|
|
print("Looking for comments...")
|
|
|
|
comments = [
|
|
c
|
|
for c in self.client.fetch_new_comments(limit=50)
|
|
if "author_name" in c
|
|
and not c["is_bot"]
|
|
and c["author_name"] != config["username"]
|
|
and c["post_id"] != 0
|
|
]
|
|
|
|
if len(comments) == 0:
|
|
print("No comments found.")
|
|
return
|
|
|
|
random.shuffle(comments)
|
|
comments = comments[: config["num_replies"]]
|
|
for comment in comments:
|
|
self.reply(comment)
|
|
|
|
def respond_to_replies(self):
|
|
replies = self.client.fetch_new_replies()
|
|
|
|
for reply in replies:
|
|
if "author_name" not in reply:
|
|
continue
|
|
|
|
if not reply["is_bot"]:
|
|
self.reply(reply)
|
|
|
|
def make_forced_replies(self):
|
|
file_path = f"{config['data_dir']}/forced.txt"
|
|
if not os.path.isfile(file_path):
|
|
return
|
|
|
|
with open(file_path, "r") as f:
|
|
lines = f.read().splitlines()
|
|
|
|
for comment_id in lines:
|
|
comment = self.client.get(f"/comment/{comment_id}")
|
|
self.reply(comment)
|
|
|
|
os.remove(file_path)
|
|
|
|
def reply(self, comment):
|
|
print(f"Generating reply for https://rdrama.net/comment/{comment['id']}")
|
|
|
|
post, thread_comments = self.client.fetch_context(comment)
|
|
|
|
if not post or not thread_comments:
|
|
print("Could not fetch context!")
|
|
|
|
prompt = utils.build_prompt(post, thread_comments)
|
|
utils.log_prompt(prompt)
|
|
|
|
candidates = []
|
|
num_candidates = config["num_candidates"] if random.random() < 0.6 else 1
|
|
while len(candidates) < num_candidates:
|
|
gen_text = self.model.generate(prompt)
|
|
reply = utils.extract_reply(gen_text)
|
|
print(f"Generated text: {gen_text}\nReply:\n{reply}")
|
|
reply = utils.format_reply(reply)
|
|
|
|
if len(reply) == 0:
|
|
print("Retrying: reply empty after processing.")
|
|
elif utils.is_low_quality(reply, post, thread_comments):
|
|
print("Retrying: low quality reply.")
|
|
else:
|
|
candidates.append(reply)
|
|
print("Accepting reply.")
|
|
|
|
reply = max(candidates, key=lambda r: min(utils.reply_length(r, 500)))
|
|
self.client.reply(reply, comment)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
bot = Bot()
|
|
bot.make_forced_replies()
|
|
bot.respond_to_replies()
|
|
bot.post_random_reply()
|
|
bot.respond_to_replies()
|