Add listener class

listener.py

@@ -0,0 +1,225 @@
+from dataclasses import dataclass
+from typing import Optional, List
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig
+import regex as re
+import torch
+import torch.nn.functional as F
+
+PROGRAM_SPECIAL_TOKEN="<extra_id_124>"
+UTTERANCES_SPECIAL_TOKEN="<extra_id_123>"
+GT_PROGRAM_SPECIAL_TOKEN="<extra_id_122>"
+
+def consistent(rx, spec):
+    # spec is in the form of (string, '+'/'-') pairs
+    for s, label in spec:
+        if not label in ['+', '-']:
+            return None
+        try:
+            if re.fullmatch(rx, s, timeout=1):
+                if label == '-':
+                    return False
+            else:
+                if label == '+':
+                    return False
+        except re.error:
+            return None
+        except TimeoutError:
+            return None
+
+    return True
+
+def get_utterance_processing_functions(label_pos, idx, separator=' '):
+    if label_pos == "suffix":
+        if idx:
+            def utterances_to_string(spec):
+                return ''.join([f"<extra_id_{i}>{s}{label}" for i, (s, label) in enumerate(spec)])
+        else:
+            def utterances_to_string(spec):
+                return separator.join([f"{s}{label}" for s, label in spec])
+    else:
+        if idx:
+            def utterances_to_string(spec):
+                return ''.join([f"<extra_id_{i}>{label}{s}" for i, (s, label) in enumerate(spec)])
+        else:
+            def utterances_to_string(spec):
+                return separator.join([f"{label}{s}" for s, label in spec])
+    
+    if label_pos == "suffix":
+        if idx:
+            def string_to_utterances(string):
+                string = re.sub(r'<extra_id_\d+>', ' ', string)
+                return [(s[:-1], s[-1]) for s in string.split(' ') if len(s) > 0]
+        else:
+            def string_to_utterances(string):
+                return [(s[:-1], s[-1]) for s in string.split(separator) if len(s) > 0]
+    else:
+        if idx:
+            def string_to_utterances(string):
+                string = re.sub(r'<extra_id_\d+>', '', string)
+                return [(s[1:], s[0]) for s in string.split(separator) if len(s) > 0]
+        else:
+            def string_to_utterances(string):
+                return [(s[1:], s[0]) for s in string.split(separator) if len(s) > 0]
+    
+    return utterances_to_string, string_to_utterances
+
+def decode(c):
+    if c < 3:
+        return f"<{c}>"
+    elif c < 258:
+        return chr(c - 3)
+    else:
+        return f"<extra_id_{c - 259}>"
+    
+def byt5_decode_batch(outputs, skip_special_tokens=True, skip_position_token=False):
+    skipped_tokens = outputs
+    if skip_special_tokens:
+        skipped_tokens = [
+            [[t for t in x if t >= 3] for x in beam]
+            for beam in skipped_tokens
+            ]
+    
+    if skip_position_token:
+        skipped_tokens = [
+            [[t for t in x if t <= 258] for x in beam] 
+            for beam in skipped_tokens
+            ]
+
+    return [
+        [''.join([decode(t) for t in x]) for x in beam]
+        for beam in skipped_tokens
+    ]
+
+class Agent:
+    def __init__(self, 
+                model_path: str,
+                gen_config: dict, 
+                device: str = "cuda", 
+                ):
+        self.device = device
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to(device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.gen_config = GenerationConfig(**gen_config)
+
+@dataclass
+class ListenerOutput:
+    programs: List[List[str]]
+    idx: Optional[List[List[int]]] = None
+    decoded: Optional[List[List[str]]] = None
+    decoded_scores: Optional[List[List[float]]] = None
+    pruned: Optional[List[List[str]]] = None
+
+
+class Listener(Agent):
+    def __init__(self, 
+        model_path, 
+        gen_config,
+        device="cuda",
+        label_pos="suffix",
+        idx: bool=True,
+        program_special_token=PROGRAM_SPECIAL_TOKEN,
+        utterances_special_token=UTTERANCES_SPECIAL_TOKEN
+    ):
+        super().__init__(
+            model_path, 
+            gen_config,
+            device=device
+        )
+        self.label_pos = label_pos
+        self.idx = idx
+        self.program_special_token = program_special_token
+        self.utterances_special_token = utterances_special_token
+        self.utterances_to_string, self.string_to_utterances = (
+            get_utterance_processing_functions(
+                label_pos, idx, separator=utterances_special_token
+                )
+            )
+    
+    def synthesize(self, context, return_scores=False, enforce_consistency=True):
+        # If context is a list of utterances, convert to string
+        if isinstance(context[0], list):
+            context_str = list(map(self.utterances_to_string, context))
+        else:
+            context_str = context
+
+        context_tokens = self.tokenizer(
+            [f"{self.utterances_special_token}{c}" if not c.startswith(self.utterances_special_token) else c 
+            for c in context_str], 
+            return_tensors="pt",
+            padding=True
+            ).to(self.device)
+        
+        decoder_inputs = self.tokenizer(
+            [self.program_special_token for _ in context], return_tensors="pt",
+            add_special_tokens=False
+            ).to(self.device)
+
+        outputs = self.model.generate(**context_tokens, 
+                                      decoder_input_ids=decoder_inputs.input_ids,
+                                      generation_config=self.gen_config, 
+                                      return_dict_in_generate=True, 
+                                      output_scores=True
+                                      )
+
+        decoded_batch = byt5_decode_batch(outputs.sequences.reshape((len(context), -1, outputs.sequences.shape[-1])).tolist(), skip_position_token=True, skip_special_tokens=True)
+
+        consistent_programs = []
+        idxs = []
+        for decoded, ctx in zip(decoded_batch, context):
+            cp = []
+            idx = []
+            for i, p in enumerate(decoded):
+                if enforce_consistency:
+                    if consistent(p, ctx):
+                        cp.append(p)
+                        idx.append(i)
+                else:
+                    cp.append(p)
+                    idx.append(i)
+            
+            consistent_programs.append(cp)
+            idxs.append(idx)
+        
+        logprobs = torch.stack(outputs.scores, dim=1).log_softmax(dim=-1)
+        gen_probs = torch.gather(logprobs, 2, outputs.sequences[:, 1:, None]).squeeze(-1)
+        gen_probs.masked_fill_(gen_probs.isinf(), 0)
+        scores = gen_probs.sum(-1)
+        n_decoded = scores.shape[0]
+        n_seq = n_decoded // len(context)
+        scores = scores.reshape((len(context), n_seq))
+        scores_list = scores.tolist()
+
+        if return_scores:
+            return ListenerOutput(
+                consistent_programs,
+                idxs, 
+                decoded_batch, 
+                scores_list
+                )
+        else:
+            return ListenerOutput(consistent_programs)
+
+    
+    def score_program(self, contexts, programs):
+        if isinstance(contexts[0], list):
+            context_str = list(map(self.utterances_to_string, contexts))
+        else:
+            context_str = contexts
+
+        context_tokens = self.tokenizer(
+            [f"{self.utterances_special_token}{c}" if not c.startswith(self.utterances_special_token) else c 
+            for c in context_str], 
+            return_tensors="pt",
+            padding=True
+            ).to(self.device)
+
+        program_tokens = self.tokenizer([f"{self.program_special_token}{p}" for p in programs], return_tensors="pt").to(self.device)
+        outputs = self.model(input_ids=context_tokens.input_ids, decoder_input_ids=program_tokens.input_ids, return_dict=True)
+        
+        logprobs = torch.gather(F.log_softmax(outputs.logits, dim=-1), 2, program_tokens.input_ids[:, 1:, None]).squeeze(-1)
+        
+        logprobs.masked_fill_(program_tokens.input_ids[:, 1:] == 0, 0)
+
+        scores = logprobs.sum(-1)
+        
+        return scores.tolist()