From 3bc9e434ba1966300f37202b74ff336f796679c9 Mon Sep 17 00:00:00 2001
From: Anthony Wang
Date: Mon, 28 Feb 2022 22:23:05 -0600
Subject: Add Matrix support 🎉

---
 README.md |  4 ++--
 bot.py    | 20 +++++++++++++++++---
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index f7f5725..d69edad 100644
--- a/README.md
+++ b/README.md
@@ -5,9 +5,9 @@ Fediverse ebooks bot using neural networks
 
 ## Usage
 
-First, install Python dependencies using your distro's package manager or `pip`: [psycopg2](https://www.psycopg.org), [torch](https://pytorch.org/), [transformers](https://huggingface.co/docs/transformers/index), and [datasets](https://huggingface.co/docs/datasets/). Additionally, for Mastodon and Pleroma, install [Mastodon.py](https://mastodonpy.readthedocs.io/en/stable/), and for Misskey, install [Misskey.py](https://misskeypy.readthedocs.io/ja/latest/). If your database or platform isn't supported, don't worry! It's easy to add support for other platforms and databases, and contributions are welcome!
+First, install Python dependencies using your distro's package manager or `pip`: [psycopg2](https://www.psycopg.org), [torch](https://pytorch.org/), [transformers](https://huggingface.co/docs/transformers/index), and [datasets](https://huggingface.co/docs/datasets/). Additionally, for Mastodon and Pleroma, install [Mastodon.py](https://mastodonpy.readthedocs.io/en/stable/), for Misskey, install [Misskey.py](https://misskeypy.readthedocs.io/ja/latest/), and for Matrix, install [simplematrixbotlib](https://simple-matrix-bot-lib.readthedocs.io/en/latest/index.html). If your database or platform isn't supported, don't worry! It's easy to add support for other platforms and databases, and contributions are welcome!
 
-Now generate the training data from your fediverse server's database using `python data.py -d 'dbname=test user=postgres password=secret host=localhost port=5432'`. You can skip this step if you have collected training data from another source.
+Now generate the training data from your fediverse server's database using `python data.py -d 'dbname=test user=postgres password=secret host=localhost port=5432'`. Generating the training data from the database is not yet supported for Matrix. You can skip this step if you have collected training data from another source.
 
 Next, train the network with `python train.py`, which may take several hours. It's a lot faster when using a GPU. If you need advanced features when training, you can also train using [run_clm.py](https://github.com/huggingface/transformers/blob/master/examples/pytorch/language-modeling/run_clm.py).
 
diff --git a/bot.py b/bot.py
index a9d661f..54cc31c 100644
--- a/bot.py
+++ b/bot.py
@@ -5,7 +5,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
 parser = ArgumentParser()
-parser.add_argument('-b', '--backend', choices=['mastodon', 'misskey'], default='mastodon',
+parser.add_argument('-b', '--backend', choices=['mastodon', 'misskey', 'matrix'], default='mastodon',
                     help='fediverse server type')
 parser.add_argument('-i', '--instance', help='Mastodon instance hosting the bot')
 parser.add_argument('-t', '--token', help='Mastodon application access token')
@@ -81,6 +81,7 @@ print(output)
 post = output.split('\n')[0]
 if len(post) < 200:
     post = output.split('\n')[0] + '\n' + output.split('\n')[1]
+post = post[:500]
 
 
 # Post it!
@@ -91,9 +92,22 @@ if args.backend == 'mastodon':
         access_token=args.token,
         api_base_url=args.instance
     )
-    mastodon.status_post(post[:500])
+    mastodon.status_post(post)
+
 elif args.backend == 'misskey':
     from Misskey import Misskey
 
     misskey = Misskey(args.instance, i=args.token)
-    misskey.notes_create(post[:500])
+    misskey.notes_create(post)
+
+elif args.backend == 'matrix':
+    import simplematrixbotlib as botlib
+
+    creds = botlib.Creds(args.instance, 'ebooks', args.token)
+    bot = botlib.Bot(creds)
+
+    @bot.listener.on_startup
+    async def room_joined(room_id):
+        await bot.api.send_text_message(room_id=room_id, message=post)
+
+    bot.run()
-- 
cgit v1.2.3-70-g09d2