Commit: 5b55807030216ac952c2c35395ff243cba572994 - zirpu (git)

Emergent generative agents

Commit MetaInfo

• Revision: 5b55807030216ac952c2c35395ff243cba572994 (tree)
• Zeit: 2023-05-20 04:04:24
• Autor: Corbin <cds@corb...>
• Commiter: Corbin

Log Message

Prototype choices, create subtagging for messages.

These go hand-in-hand; we want to be able to choose to not just reply on
IRC, but reply on a specific channel, so we need to use the channel as
subtag data.

The choice mechanism is clunky; I would very much like to just use
ReLLM, and maybe that'll be important in the future if/when we migrate
back to HF with their brand-new RWKV integration. They don't support
quantized models, though...

Ändern Zusammenfassung

• modified: src/agent.py
• modified: src/gens/mawrkov.py

Diff

--- a/src/agent.py
+++ b/src/agent.py
@@ -6,18 +6,19 @@ import json
 import os.path
 import re
 import random
+from string import ascii_uppercase
 import sys
 from threading import Lock
 
 from twisted.internet import reactor
 from twisted.internet.defer import DeferredLock, succeed
 from twisted.internet.protocol import ClientFactory
-from twisted.internet.task import LoopingCall
+from twisted.internet.task import LoopingCall, deferLater
 from twisted.internet.threads import deferToThread
 from twisted.words.protocols.irc import IRCClient
 
 from common import irc_line, Timer, SentenceIndex, breakAt
-from gens.mawrkov import MawrkovGen
+from gens.mawrkov import MawrkovGen, force
 from gens.trans import SentenceEmbed
 
 build_traits = " + ".join
@@ -43,55 +44,85 @@ prologues = {
     "clock": "Then I checked the time:",
     "lbi": "Then I tried to interpret what just happened:",
     "thoughts": "Then I thought to myself:",
-    "irc": "Then I chatted on IRC:",
+    "choice": "Then I chose what to do next:",
+    "irc": "Then I chatted on IRC, channel {subtag}:",
 }
 
 class Mind:
     currentTag = None
     logits = state = None
 
-    def __init__(self): self.lock = DeferredLock()
+    def __init__(self, name):
+        # Exclusively for debugging!
+        self.name = name
+        self.lock = DeferredLock()
 
-    def switchTag(self, tag):
-        if tag == self.currentTag: return succeed(None)
+    def switchTag(self, tag, subtag):
+        if (tag, subtag) == self.currentTag: return succeed(None)
         else:
-            self.currentTag = tag
+            self.currentTag = tag, subtag
             # Double newlines are added here.
-            return deferToThread(self.write, "\n" + prologues[tag])
+            return deferToThread(self.write, "\n" +
+                                 prologues[tag].format(subtag=subtag))
 
-    def overhear(self, tag, s):
+    def overhear(self, tag, subtag, s):
         def cb():
-            d = self.switchTag(tag)
+            d = self.switchTag(tag, subtag)
             d.addCallback(lambda _: deferToThread(self.write, s))
             return d
         return self.lock.run(cb)
 
-    def write(self, s):
+    def writeRaw(self, s):
         with genLock:
-            # Newlines are added here.
-            self.logits, self.state = gen.feedForward(gen.tokenize(s + "\n"),
+            self.logits, self.state = gen.feedForward(gen.tokenize(s),
                                                       self.logits, self.state)
-            print(s)
+            print(self.name, "~~", s.strip())
+
+    # Newlines are added here.
+    def write(self, s): self.writeRaw(s + "\n")
 
     def complete(self, s):
         with genLock:
             completion, self.logits, self.state = gen.complete(s, self.logits, self.state)
-            print("«", completion, "»")
+            print(self.name, "«", completion.strip(), "»")
             return completion
 
-    def infer(self, tag, prefix):
+    def infer(self, tag, subtag, prefix):
+        def cb():
+            d = self.switchTag(tag, subtag)
+            d.addCallback(lambda _: print(self.name, "~?", prefix)
+                          or deferToThread(self.complete, prefix))
+            return d
+        return self.lock.run(cb)
+
+    def forceChoice(self, tag, subtag, prefix, options):
+        tokens = [gen.tokenize(opt)[0] for opt in options]
         def cb():
-            print(prefix)
-            d = self.switchTag(tag)
-            d.addCallback(lambda _: deferToThread(self.complete, prefix))
+            d = self.switchTag(tag, subtag)
+            # Breaking the newline invariant...
+            d.addCallback(lambda _: print(self.name, "~!", prefix)
+                          or deferToThread(self.writeRaw, prefix))
+            # ...so that this inference happens before the newline...
+            d.addCallback(lambda _: force(tokens, self.logits))
+            d.addCallback(lambda t: gen.tokenizer.decode([t]))
+
+            @d.addCallback
+            def cb2(s):
+                # ...and now the invariant must be restored...
+                d2 = deferToThread(self.write, s)
+                # ...but we want to tell the caller what was chosen!
+                d2.addCallback(lambda _: s)
+                return d2
             return d
         return self.lock.run(cb)
 
 
 class Agent:
     listeners = ()
+    subtag = None
     def broadcast(self, s):
-        for listener in self.listeners: listener.overhear(self.tag, s)
+        for listener in self.listeners:
+            listener.overhear(self.tag, self.subtag, s)
 
 class Clock(Agent):
     tag = "clock"
@@ -102,45 +133,77 @@ class LeftBrainInterpreter(Agent):
     tag = "lbi"
     events = 0
     def __init__(self, mind): self.mind = mind
-    def overhear(self, tag, s):
+    def overhear(self, tag, subtag, s):
+        # XXX should skip thoughts, maybe?
         self.events += 1
-        if self.events >= 10:
+        if self.events >= 20:
             self.events = 0
-            return self.mind.infer(self.tag, "")
+            return self.mind.infer(self.tag, self.subtag, "")
+
+choices = {
+    "irc": "Chat on IRC, channel {subtag}",
+    "thoughts": "Think to myself for a moment",
+}
+def choicify(options):
+    return " ".join(f"({c}) {o}" for c, o in zip(ascii_uppercase, options))
+
+class ChoiceMaker(Agent):
+    tag = "choice"
+    def __init__(self, mind):
+        self.mind = mind
+        self.possibilities = set()
+        self.actions = {}
+    def overhear(self, tag, subtag, s): self.possibilities.add((tag, subtag))
+    def dispatch(self, tag, subtag):
+        if tag in self.actions: return self.actions[tag](subtag)
+    def idle(self):
+        possibilities = [(tag, subtag)
+                         for tag, subtag in self.possibilities
+                         if tag in choices]
+        if len(possibilities) == 0: return
+        elif len(possibilities) == 1: return self.dispatch(*possibilities[0])
+
+        options = [choices[tag].format(subtag=subtag)
+                   for tag, subtag in possibilities]
+        prompt = choicify(options.values())
+        self.broadcast(prompt)
+        d = self.mind.forceChoice(self.tag, self.subtag, "My choice: ",
+                                  ascii_uppercase[:len(options)])
+        d.addCallback(lambda s: self.dispatch(*possibilities[chr(ord(s[0]) -
+                                                                 ord('A'))]))
+        return d
 
 class ChainOfThoughts(Agent):
     tag = "thoughts"
-    def __init__(self, mind, index, seed):
+    def __init__(self, mind, index, seed, idle):
         self.mind = mind
         self.index = index
         self.recentThoughts = deque([seed], maxlen=5)
+        self.idle = idle
 
-    def go(self):
-        cb = self.reflect if random.choice([0, 1]) else self.cogitate
-        return cb()
+    def go(self): return random.choice([self.reflect, self.cogitate])()
 
-    def overhear(self, tag, s):
-        self.addRelatedThoughts(tag + ": " + s)
+    def overhear(self, tag, subtag, s):
+        prefix = f"{tag} ({subtag}): " if subtag else tag + ": "
+        self.addRelatedThoughts(prefix + s)
 
     def addRelatedThoughts(self, s):
         thoughts = self.index.search(s, 2)
+        new = 0
         for thought in thoughts:
             if thought not in self.recentThoughts:
+                new += 1
                 self.recentThoughts.append(thought)
                 self.broadcast(thought)
+        return new
 
-    def cogitate(self): self.addRelatedThoughts(self.recentThoughts[-1])
+    def cogitate(self):
+        new = self.addRelatedThoughts(self.recentThoughts[-1])
+        if not new: return deferLater(reactor, 0.0, self.idle)
     def reflect(self):
-        d = self.mind.infer(self.tag, "")
-
-        @d.addCallback
-        def cb(s):
-            if not s.strip():
-                self.broadcast(random.choice([
-                    "Head empty; no thoughts.",
-                    "So bored.",
-                    "Zoned out.",
-                ]))
+        d = self.mind.infer(self.tag, self.subtag, "")
+        d.addCallback(lambda s: s.strip() or self.idle())
+        return d
 
 
 IRC_LINE_HEAD = re.compile(r"\d{1,2}:\d{1,2}:\d{1,2} <")
@@ -155,40 +218,45 @@ class IRCAgent(Agent, IRCClient):
         self.nickname = title_to_nick(title)
         self.startingChannels = startingChannels
 
+    def broadcastTagged(self, subtag, s):
+        for listener in self.listeners: listener.overhear(self.tag, subtag, s)
+
     def prefix(self, channel):
         return f"{datetime.now():%H:%M:%S} {channel} <{self.nickname}>"
 
     def userJoined(self, user, channel):
-        self.broadcast(f"{user} joins {channel}")
+        self.broadcastTagged(channel, f"{user} joins {channel}")
 
     def userLeft(self, user, channel):
-        self.broadcast(f"{user} leaves {channel}")
+        self.broadcastTagged(channel, f"{user} leaves {channel}")
 
     def userQuit(self, user, channel):
-        self.broadcast(f"{user} quits {channel}")
+        self.broadcastTagged(channel, f"{user} quits {channel}")
 
     def topicUpdated(self, user, channel, topic):
-        self.broadcast(f"Topic for {channel} is now: {topic}")
+        self.broadcastTagged(channel, f"Topic for {channel} is now: {topic}")
 
     def signedOn(self):
         for channel in self.startingChannels: self.join(channel)
 
     def privmsg(self, user, channel, line):
         user = user.split("!", 1)[0]
-        self.broadcast(irc_line(datetime.now(), channel, user, line))
-        if self.nickname in line:
-            d = self.mind.infer("irc", self.prefix(channel))
+        self.broadcastTagged(channel, irc_line(datetime.now(), channel, user, line))
+        # XXX should eventually be not needed!
+        # if self.nickname in line: self.speakInChannel(channel)
 
-            @d.addCallback
-            def cb(s):
-                line = breakIRCLine(s).strip()
-                self.msg(channel, line)
+    def speakInChannel(self, channel):
+        d = self.mind.infer("irc", channel, self.prefix(channel))
+
+        @d.addCallback
+        def cb(s): self.msg(channel, breakIRCLine(s).strip())
 
 class IRCFactory(ClientFactory):
     protocol = IRCAgent
-    def __init__(self, mind, listeners, title, startingChannels):
+    def __init__(self, mind, cm, listeners, title, startingChannels):
         super(IRCFactory, self).__init__()
         self.mind = mind
+        self.cm = cm
         self.listeners = listeners
         self.title = title
         self.startingChannels = startingChannels
@@ -196,12 +264,13 @@ class IRCFactory(ClientFactory):
         protocol = self.protocol(self.mind, self.title, self.startingChannels)
         protocol.factory = self
         protocol.listeners = self.listeners
+        self.cm.actions["irc"] = protocol.speakInChannel
         return protocol
 
 def go():
     print("~ Starting tasks…")
     clock = Clock()
-    LoopingCall(clock.go).start(300.0, now=False)
+    LoopingCall(clock.go).start(60 * 10, now=False)
 
     for logpath in sys.argv[2:]:
         character = load_character(logpath)
@@ -210,19 +279,21 @@ def go():
         thoughtPath = os.path.join(logpath, "thoughts.txt")
         thoughtIndex = SentenceIndex.fromPath(thoughtPath, embedder)
 
-        mind = Mind()
+        mind = Mind(title_to_nick(title))
         with Timer("initial warmup"):
             mind.logits, mind.state = gen.feedForward(gen.tokenize(firstStatement), None, None)
 
         lbi = LeftBrainInterpreter(mind)
         clock.listeners += lbi, mind
 
-        thoughts = ChainOfThoughts(mind, thoughtIndex, firstStatement)
+        cm = ChoiceMaker(mind)
+
+        thoughts = ChainOfThoughts(mind, thoughtIndex, firstStatement, cm.idle)
         thoughts.listeners = lbi, mind
-        LoopingCall(thoughts.go).start(120.0, now=False)
+        LoopingCall(thoughts.go).start(60, now=False)
 
         print("~ Thought index:", thoughtIndex.size(), "thoughts")
-        factory = IRCFactory(mind, (thoughts, lbi, mind),
+        factory = IRCFactory(mind, cm, (thoughts, lbi, cm, mind),
                              title,
                              character["startingChannels"])
         print("~ Connecting factory for:", title)

--- a/src/gens/mawrkov.py
+++ b/src/gens/mawrkov.py
@@ -27,8 +27,7 @@ TEMPERATURE = 0.8
 TOP_P = 0.8
 
 class MawrkovGen:
-    # XXX might be wrong
-    model_name = "The Pile (14B params, 4-bit quantized)"
+    model_name = "The Pile"
     model_arch = "RWKV"
     def __init__(self, model_path, max_new_tokens):
         self.max_new_tokens = max_new_tokens
@@ -57,3 +56,8 @@ class MawrkovGen:
             logits, state = self.feedForward([token], logits, state)
             if "\n" in self.tokenizer.decode([token]): break
         return self.tokenizer.decode(tokens).split("\n", 1)[0], logits, state
+
+def force(options, logits):
+    # +10 is hopefully not too much.
+    biases = {opt: 10 for opt in options}
+    return sampling.sample_logits(logits, TEMPERATURE, TOP_P, biases)