improve performance

app.py

@@ -1,7 +1,8 @@
 import spacy
 import nltk
 nltk.download('wordnet', quiet=True)
-spacy.cli.download('en_core_web_sm')
+if not spacy.util.is_package('en_core_web_sm'):
+    spacy.cli.download('en_core_web_sm')
 from compute_lng import compute_lng
 
 import torch
@@ -111,7 +112,7 @@ def impute_targets():
     shared_state.target = round_ling(interp_raw).tolist()
     return shared_state.target
 
-def generate_with_feedback(sent1, approx):
+def generate_with_feedback(sent1, approx, progress=gr.Progress()):
     if sent1 == '':
         raise gr.Error('Please input a source text.')
 
@@ -122,24 +123,25 @@ def generate_with_feedback(sent1, approx):
     input_ids = tokenizer.encode(sent1, return_tensors='pt').to(device)
     ling2 = torch.tensor(scaler.transform([shared_state.target])).float().to(device)
     inputs = {
-            'sentence1_input_ids': input_ids,
+            'input_ids': input_ids,
             'sentence2_ling': ling2,
-            'sentence1_attention_mask': torch.ones_like(input_ids)
+            'attention_mask': torch.ones_like(input_ids)
               }
 
-    pred, (pred_text, interpolations) = model.infer_with_feedback_BP(ling_disc, sem_emb, inputs, tokenizer)
+    progress((0, None), unit='intermediate paraphrase generated.')
+    pred, (pred_text, interpolations) = model.infer_with_feedback_BP(ling_disc, sem_emb, inputs, tokenizer, progress)
 
     interpolation = '-- ' + '\n-- '.join(interpolations)
     # Return both the generation results and the updated slider values
     return [pred_text, interpolation] + [gr.update(value=val) for val in shared_state.target]
 
-def generate_random(sent1, count, approx):
+def generate_random(sent1, count, approx, progress=gr.Progress()):
     if sent1 == '':
         raise gr.Error('Please input a source text.')
     preds, interpolations = [], []
     orig_active_indices = shared_state.active_indices
     shared_state.active_indices = set(range(len(lng_names)))
-    for c in range(count):
+    for c in progress.tqdm(range(count), desc='Generating random sentences', unit='paraphrases'):
         idx = np.random.randint(0, len(ling_collection))
         ling_ex = ling_collection[idx]
         shared_state.target = ling_ex.copy()
@@ -167,7 +169,7 @@ def generate_random(sent1, count, approx):
     shared_state.active_indices = orig_active_indices
     return '\n***\n'.join(preds), '\n***\n'.join(interpolations)
 
-def estimate_gen(sent1, sent2, approx):
+def estimate_gen(sent1, sent2, approx, progress=gr.Progress()):
     if 'approximate' in approx:
         input_ids = tokenizer.encode(sent2, return_tensors='pt').to(device)
         with torch.no_grad():
@@ -183,7 +185,7 @@ def estimate_gen(sent1, sent2, approx):
     
     orig_active_indices = shared_state.active_indices
     shared_state.active_indices = set(range(len(lng_names)))
-    gen = generate_with_feedback(sent1, approx)[:2]
+    gen = generate_with_feedback(sent1, approx, progress)[:2]
     shared_state.active_indices = orig_active_indices
     return gen + [gr.update(value=val) for val in shared_state.target]
 

const.py

@@ -1030,7 +1030,6 @@ used_indices = [
         63, 64, 65, 66, 67, 68, 73, 121, 124, 129, 134, 136, 254,
         257, 258, 261, 263, 272, 274 
         ]
-lftk_used_indices = [1, 7, 8, 9, 10, 11, 12, 17, 65, 68, 73, 78, 80, 198, 201, 202, 205, 207, 216, 218]
 
 eval_indices = [4,5,6,18,257,272]
 eval_indices = [used_indices.index(idx) for idx in eval_indices]

model.py

@@ -10,6 +10,10 @@ from types import MethodType
 from utils import *
 from ling_disc import DebertaReplacedTokenizer
 from const import *
+from lingconv_t5 import LingConvT5ForConditionalGeneration
+from dataclasses import dataclass
+from transformers.modeling_outputs import Seq2SeqLMOutput
+from typing import Optional, Dict, Any
 
 
 
@@ -77,9 +81,9 @@ class LingGenerator(nn.Module):
         bs = inputs_embeds.shape[0]
 
         if self.gen_input == 's+l':
-            sent1_ling = self.ling_embed(batch['sentence1_ling'])
-            sent1_ling = sent1_ling.view(bs, 1, -1)
-            inputs_embeds = inputs_embeds + sent1_ling
+            sentence1_ling = self.ling_embed(batch['sentence1_ling'])
+            sentence1_ling = sentence1_ling.view(bs, 1, -1)
+            inputs_embeds = inputs_embeds + sentence1_ling
 
         gen = self.gen(inputs_embeds=inputs_embeds,
                 attention_mask=inputs_att_mask).last_hidden_state.mean(1)
@@ -185,13 +189,13 @@ class SemEmb(T5EncoderModel):
                 nn.Linear(hidden_dim, 1))
 
     def compare_sem(self, **batch):
-        bs = batch['sentence1_attention_mask'].shape[0]
-        ones = torch.ones((bs, 1), device=batch['sentence1_attention_mask'].device)
+        bs = batch['attention_mask'].shape[0]
+        ones = torch.ones((bs, 1), device=batch['attention_mask'].device)
         sep = torch.ones((bs, 1), dtype=torch.long,
-                device=batch['sentence1_attention_mask'].device) * self.sep_token_id
-        att_mask = torch.cat([batch['sentence1_attention_mask'], ones, batch['sentence2_attention_mask']], dim=1)
+                device=batch['attention_mask'].device) * self.sep_token_id
+        att_mask = torch.cat([batch['attention_mask'], ones, batch['sentence2_attention_mask']], dim=1)
         if 'logits' in batch:
-            input_ids = torch.cat([batch['sentence1_input_ids'], sep], dim=1)
+            input_ids = torch.cat([batch['input_ids'], sep], dim=1)
             embeds1 = self.shared(input_ids)
 
             logits = batch['logits']
@@ -201,11 +205,11 @@ class SemEmb(T5EncoderModel):
 
             embeds2 =  onehot_ @ self.shared.weight
             embeds1_2 = torch.cat([embeds1, embeds2], dim=1)
-            hidden_units = self(inputs_embeds=embeds1_2,
+            hidden_units = super().forward(inputs_embeds=embeds1_2,
                     attention_mask=att_mask).last_hidden_state.mean(1)
         elif 'sentence2_input_ids' in batch:
-            input_ids = torch.cat([batch['sentence1_input_ids'], sep, batch['sentence2_input_ids']], dim=1)
-            hidden_units = self(input_ids=input_ids,
+            input_ids = torch.cat([batch['input_ids'], sep, batch['sentence2_input_ids']], dim=1)
+            hidden_units = super().forward(input_ids=input_ids,
                     attention_mask=att_mask).last_hidden_state.mean(1)
         probs = self.projection(hidden_units)
         return probs
@@ -222,31 +226,36 @@ def prepare_inputs_for_generation(
         cross_attn_head_mask=None,
         use_cache=None,
         encoder_outputs=None,
-        sent1_ling=None,
-        sent2_ling=None,
+        sentence1_ling=None,
+        sentence2_ling=None,
         **kwargs
     ):
-
         # cut decoder_input_ids if past is used
         if past_key_values is not None:
             input_ids = input_ids[:, -1:]
 
+        cached = use_cache and len(past_key_values) > 0
+
         input_ids = input_ids.clone()
         decoder_inputs_embeds = self.shared(input_ids)
 
-        if combine_method == 'decoder_add_first':
-            sent2_ling = torch.cat([sent2_ling,
-                torch.repeat_interleave(torch.zeros_like(sent2_ling), input_ids.shape[1] - 1, dim=1)], dim = 1)
-        if combine_method == 'decoder_concat':
+        if combine_method == 'layer_injection':
+            # For layer injection, we'll pass the ling embeddings separately
+            ling_embed = sentence2_ling if ling2_only else (sentence1_ling + sentence2_ling)
+        elif combine_method == 'decoder_add_first' and not cached:
+            sentence2_ling = torch.cat([sentence2_ling,
+                torch.repeat_interleave(torch.zeros_like(sentence2_ling), input_ids.shape[1] - 1, dim=1)], dim = 1)
+        elif combine_method == 'decoder_concat':
             if ling2_only:
-                decoder_inputs_embeds = torch.cat([sent2_ling, decoder_inputs_embeds], dim=1)
+                decoder_inputs_embeds = torch.cat([sentence2_ling, decoder_inputs_embeds], dim=1)
             else:
-                decoder_inputs_embeds = torch.cat([sent1_ling, sent2_ling, decoder_inputs_embeds], dim=1)
-        elif combine_method == 'decoder_add'or (past_key_values is None and combine_method == 'decoder_add_first'):
+                decoder_inputs_embeds = torch.cat([sentence1_ling, sentence2_ling, decoder_inputs_embeds], dim=1)
+
+        if combine_method == 'decoder_add' or (not cached and combine_method == 'decoder_add_first'):
             if ling2_only:
-                decoder_inputs_embeds = decoder_inputs_embeds + sent2_ling
+                decoder_inputs_embeds = decoder_inputs_embeds + sentence2_ling
             else:
-                decoder_inputs_embeds = decoder_inputs_embeds + sent1_ling + sent2_ling
+                decoder_inputs_embeds = decoder_inputs_embeds + sentence1_ling + sentence2_ling
 
         return {
             "decoder_inputs_embeds": decoder_inputs_embeds,
@@ -257,19 +266,27 @@ def prepare_inputs_for_generation(
             "decoder_head_mask": decoder_head_mask,
             "cross_attn_head_mask": cross_attn_head_mask,
             "use_cache": use_cache,
+            "ling_embed": ling_embed if combine_method == 'layer_injection' else None,
         }
 
 class LogitsAdd(LogitsProcessor):
-    def __init__(self, sent2_ling):
+    def __init__(self, sentence2_ling):
         super().__init__()
-        self.sent2_ling = sent2_ling
+        self.sentence2_ling = sentence2_ling
 
     def __call__(self, input_ids, scores):
-        return scores + self.sent2_ling
+        return scores + self.sentence2_ling
 
-class EncoderDecoderVAE(T5ForConditionalGeneration):
+class EncoderDecoderVAE(LingConvT5ForConditionalGeneration):
     def __init__(self, config, args, pad_token_id, sepeos_token_id, vocab_size = 32128):
+        if args.combine_method == 'layer_injection':
+            if args.injection_layer < 0 or args.injection_layer >= config.num_decoder_layers:
+                raise ValueError(f"Invalid injection layer: {args.injection_layer}. Must be between 0 and {config.num_decoder_layers - 1}.")
+            config.ling_injection_layer = args.injection_layer
+            config.ling_injection_type = args.injection_type  # 'first' or 'all'
+            
         super().__init__(config)
+        
         self.prepare_inputs_for_generation = types.MethodType(
                 partial(prepare_inputs_for_generation, args.combine_method, args.ling2_only),
                 self)
@@ -287,7 +304,7 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
                     nn.ReLU(),
                     nn.Linear(hidden_dim, hidden_dim),
                     )
-        elif 'concat' in args.combine_method or 'add' in args.combine_method:
+        elif 'concat' in args.combine_method or 'add' in args.combine_method or 'layer_injection' in args.combine_method:
             if args.ling_embed_type == 'two-layer':
                 self.ling_embed = nn.Sequential(
                         nn.Linear(args.lng_dim, args.lng_dim),
@@ -297,6 +314,7 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
             else:
                 self.ling_embed = nn.Linear(args.lng_dim, hidden_dim)
             self.ling_dropout = nn.Dropout(args.ling_dropout)
+        self.ling_embed.apply(self._init_weights)
 
         if args.ling_vae:
             self.ling_mu = nn.Linear(hidden_dim, hidden_dim)
@@ -306,8 +324,20 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
             nn.init.xavier_uniform_(self.ling_logvar.weight)
 
 
-        generate_with_grad = unwrap(self.generate)
+        generate_with_grad = unwrap(super().generate)
         self.generate_with_grad = MethodType(generate_with_grad, self)
+        self.generate_original = super().generate
+
+    def _init_weights(self, module):
+        std = self.args.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
 
     def get_fusion_layer(self):
         if 'fusion' in self.args.combine_method:
@@ -321,122 +351,143 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
         std = torch.exp(0.5 * logvar)
         return mu + std * torch.randn_like(std)
 
-    def encode(self, batch):
-        if 'inputs_embeds' in batch:
-            inputs_embeds = batch['inputs_embeds']
+    def _process_ling_embeddings(self, sentence1_ling, sentence2_ling, 
+                               sentence1_ling_embed, sentence2_ling_embed, bs):
+        """Helper method to process linguistic embeddings"""
+        cache = {}
+        
+        # Process sentence1 embedding
+        if sentence1_ling_embed is not None:
+            sentence1_ling = sentence1_ling_embed
+        elif sentence1_ling is not None:
+            sentence1_ling = self.ling_embed(self.ling_dropout(sentence1_ling))
         else:
-            inputs_embeds = self.shared(batch['sentence1_input_ids'])
-        inputs_att_mask = batch['sentence1_attention_mask']
+            sentence1_ling = None
+            
+        # Process sentence2 embedding
+        if sentence2_ling_embed is not None:
+            sentence2_ling = sentence2_ling_embed
+        elif sentence2_ling is not None:
+            sentence2_ling = self.ling_embed(self.ling_dropout(sentence2_ling))
+        else:
+            sentence2_ling = None
+
+        # Apply VAE if configured
+        if self.args.ling_vae and sentence1_ling is not None and sentence2_ling is not None:
+            sentence1_ling = F.leaky_relu(sentence1_ling)
+            sent1_mu, sent1_logvar = self.ling_mu(sentence1_ling), self.ling_logvar(sentence1_ling)
+            sentence1_ling = self.sample(sent1_mu, sent1_logvar)
+
+            sentence2_ling = F.leaky_relu(sentence2_ling)
+            sent2_mu, sent2_logvar = self.ling_mu(sentence2_ling), self.ling_logvar(sentence2_ling)
+            sentence2_ling = self.sample(sent2_mu, sent2_logvar)
+            
+            cache.update({
+                'sent1_mu': sent1_mu, 'sent1_logvar': sent1_logvar,
+                'sent2_mu': sent2_mu, 'sent2_logvar': sent2_logvar,
+                'sentence1_ling': sentence1_ling, 'sentence2_ling': sentence2_ling
+            })
+        else:
+            if sentence2_ling is not None:
+                cache['sentence2_ling'] = sentence2_ling
+            if sentence1_ling is not None:
+                cache['sentence1_ling'] = sentence1_ling
+
+        # Reshape embeddings
+        if sentence1_ling is not None:
+            sentence1_ling = sentence1_ling.view(bs, 1, -1)
+        if sentence2_ling is not None:
+            sentence2_ling = sentence2_ling.view(bs, 1, -1)
+            
+        return sentence1_ling, sentence2_ling, cache
+
+    def encode(self, 
+               input_ids=None,
+               attention_mask=None,
+               sentence1_ling=None,
+               sentence2_ling=None,
+               sentence1_ling_embed=None,
+               sentence2_ling_embed=None,
+               inputs_embeds=None,
+               ):
+        if inputs_embeds is None:
+            inputs_embeds = self.shared(input_ids)
+        inputs_att_mask = attention_mask if attention_mask is not None else torch.ones_like(input_ids)
         bs = inputs_embeds.shape[0]
-        cache = {}
+        
         if self.args.combine_method in ('input_concat', 'input_add'):
-            if 'sent1_ling_embed' in batch:
-                sent1_ling = batch['sent1_ling_embed']
-            else:
-                sent1_ling = self.ling_embed(self.ling_dropout(batch['sentence1_ling']))
-            if 'sent2_ling_embed' in batch:
-                sent2_ling = batch['sent2_ling_embed']
-            else:
-                sent2_ling = self.ling_embed(self.ling_dropout(batch['sentence2_ling']))
-            if self.args.ling_vae:
-                sent1_ling = F.leaky_relu(sent1_ling)
-                sent1_mu, sent1_logvar = self.ling_mu(sent1_ling), self.ling_logvar(sent1_ling)
-                sent1_ling = self.sample(sent1_mu, sent1_logvar)
-
-                sent2_ling = F.leaky_relu(sent2_ling)
-                sent2_mu, sent2_logvar = self.ling_mu(sent2_ling), self.ling_logvar(sent2_ling)
-                sent2_ling = self.sample(sent2_mu, sent2_logvar)
-                cache.update({'sent1_mu': sent1_mu, 'sent1_logvar': sent1_logvar,
-                    'sent2_mu': sent2_mu, 'sent2_logvar': sent2_logvar,
-                    'sent1_ling': sent1_ling, 'sent2_ling': sent2_ling})
-            else:
-                cache.update({'sent1_ling': sent1_ling, 'sent2_ling': sent2_ling})
-            sent1_ling = sent1_ling.view(bs, 1, -1)
-            sent2_ling = sent2_ling.view(bs, 1, -1)
+            sentence1_ling, sentence2_ling, cache = self._process_ling_embeddings(
+                sentence1_ling, sentence2_ling,
+                sentence1_ling_embed, sentence2_ling_embed, bs
+            )
+            
             if self.args.combine_method == 'input_concat':
                 if self.args.ling2_only:
-                    inputs_embeds = torch.cat([inputs_embeds, sent2_ling], dim=1)
+                    inputs_embeds = torch.cat([inputs_embeds, sentence2_ling], dim=1)
                     inputs_att_mask = torch.cat([inputs_att_mask,
                         torch.ones((bs, 1)).to(inputs_embeds.device)], dim=1)
                 else:
-                    inputs_embeds = torch.cat([inputs_embeds, sent1_ling, sent2_ling], dim=1)
+                    inputs_embeds = torch.cat([inputs_embeds, sentence1_ling, sentence2_ling], dim=1)
                     inputs_att_mask = torch.cat([inputs_att_mask,
                         torch.ones((bs, 2)).to(inputs_embeds.device)], dim=1)
             elif self.args.combine_method == 'input_add':
                 if self.args.ling2_only:
-                    inputs_embeds = inputs_embeds + sent2_ling
+                    inputs_embeds = inputs_embeds + sentence2_ling
                 else:
-                    inputs_embeds = inputs_embeds + sent1_ling + sent2_ling
+                    inputs_embeds = inputs_embeds + sentence1_ling + sentence2_ling
+        else:
+            cache = {}
+
         return self.encoder(inputs_embeds=inputs_embeds,
                 attention_mask=inputs_att_mask), inputs_att_mask, cache
 
-    def decode(self, batch, enc_output, inputs_att_mask, generate):
-        bs = inputs_att_mask.shape[0]
+    def decode(self,
+              sentence2_input_ids=None,
+              sentence1_ling=None,
+              sentence2_ling=None,
+              encoder_outputs=None,
+              encoder_attention_mask=None,
+              decoder_inputs_embeds=None,
+              decoder_attention_mask=None,
+              generate=False,
+              sentence1_ling_embed=None,
+              sentence2_ling_embed=None,
+              ling_embed=None,
+              generate_with_grad=False,
+              **kwargs
+              ):
+        bs = encoder_outputs[0].shape[0]
         cache = {}
-        if self.args.combine_method in ('embed_concat', 'decoder_concat', 'decoder_add', 'logits_add', 'decoder_add_first'):
-            if 'sent1_ling_embed' in batch:
-                sent1_ling = batch['sent1_ling_embed']
-            elif 'sentence1_ling' in batch:
-                sent1_ling = self.ling_embed(self.ling_dropout(batch['sentence1_ling']))
-            else:
-                sent1_ling = None
-            if 'sent2_ling_embed' in batch:
-                sent2_ling = batch['sent2_ling_embed']
-            else:
-                sent2_ling = self.ling_embed(self.ling_dropout(batch['sentence2_ling']))
-            if self.args.ling_vae:
-                sent1_ling = F.leaky_relu(sent1_ling)
-                sent1_mu, sent1_logvar = self.ling_mu(sent1_ling), self.ling_logvar(sent1_ling)
-                sent1_ling = self.sample(sent1_mu, sent1_logvar)
-
-                sent2_ling = F.leaky_relu(sent2_ling)
-                sent2_mu, sent2_logvar = self.ling_mu(sent2_ling), self.ling_logvar(sent2_ling)
-                sent2_ling = self.sample(sent2_mu, sent2_logvar)
-                cache.update({'sent1_mu': sent1_mu, 'sent1_logvar': sent1_logvar,
-                    'sent2_mu': sent2_mu, 'sent2_logvar': sent2_logvar,
-                    'sent1_ling': sent1_ling, 'sent2_ling': sent2_ling})
-            else:
-                cache.update({'sent2_ling': sent2_ling})
-                if sent1_ling is not None:
-                    cache.update({'sent1_ling': sent1_ling})
-            if sent1_ling is not None:
-                sent1_ling = sent1_ling.view(bs, 1, -1)
-            sent2_ling = sent2_ling.view(bs, 1, -1)
-            if self.args.combine_method == 'decoder_add_first' and not generate:
-                sent2_ling = torch.cat([sent2_ling,
-                    torch.repeat_interleave(torch.zeros_like(sent2_ling), batch['sentence2_input_ids'].shape[1] - 1, dim=1)], dim = 1)
-        else:
-            sent1_ling, sent2_ling = None, None
-
-        if self.args.combine_method == 'embed_concat':
-            enc_output.last_hidden_state = torch.cat([enc_output.last_hidden_state,
-                sent1_ling, sent2_ling], dim=1)
-            inputs_att_mask = torch.cat([inputs_att_mask,
-                torch.ones((bs, 2)).to(inputs_att_mask.device)], dim=1)
-        elif 'fusion' in self.args.combine_method:
-            sent1_ling = batch['sentence1_ling'].unsqueeze(1)\
-                    .expand(-1, enc_output.last_hidden_state.shape[1], -1)
-            sent2_ling = batch['sentence2_ling'].unsqueeze(1)\
-                    .expand(-1, enc_output.last_hidden_state.shape[1], -1)
-            if self.args.ling2_only:
-                combined_embedding = torch.cat([enc_output.last_hidden_state, sent2_ling], dim=2)
+        
+        if decoder_inputs_embeds is None:
+            if self.args.combine_method in ('embed_concat', 'decoder_concat', 'decoder_add', 
+                                        'logits_add', 'decoder_add_first', 'layer_injection'):
+                sentence1_ling, sentence2_ling, cache = self._process_ling_embeddings(
+                    sentence1_ling, sentence2_ling,
+                    sentence1_ling_embed, sentence2_ling_embed, bs
+                )
+                
+                if (self.args.combine_method == 'decoder_add_first' or 
+                    (self.args.combine_method == 'layer_injection' and 
+                    self.args.injection_type == 'first')) and not generate:
+                    sentence2_ling = torch.cat([sentence2_ling,
+                        torch.repeat_interleave(torch.zeros_like(sentence2_ling), 
+                        sentence2_input_ids.shape[1] - 1, dim=1)], dim = 1)
             else:
-                combined_embedding = torch.cat([enc_output.last_hidden_state, sent1_ling, sent2_ling], dim=2)
-            enc_output.last_hidden_state = self.fusion(combined_embedding)
-
+                sentence1_ling, sentence2_ling = None, None
+        
         if generate:
             if self.args.combine_method == 'logits_add':
-                logits_processor = LogitsProcessorList([LogitsAdd(sent2_ling.view(bs, -1))])
+                logits_processor = LogitsProcessorList([LogitsAdd(sentence2_ling.view(bs, -1))])
             else:
                 logits_processor = LogitsProcessorList()
 
-            dec_output = self.generate_with_grad(
-                    attention_mask=inputs_att_mask,
-                    encoder_outputs=enc_output,
-                    sent1_ling=sent1_ling,
-                    sent2_ling=sent2_ling,
-                    return_dict_in_generate=True,
-                    output_scores=True,
+            generate_fn = self.generate_with_grad if generate_with_grad else self.generate_original
+            dec_output = generate_fn(
+                    attention_mask=encoder_attention_mask,
+                    encoder_outputs=encoder_outputs,
+                    sentence1_ling=sentence1_ling,
+                    sentence2_ling=sentence2_ling,
                     logits_processor = logits_processor,
                     # renormalize_logits=True,
                     # do_sample=True,
@@ -445,68 +496,135 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
                     # min_new_tokens=3,
                     # repetition_penalty=1.2,
                     max_length=self.args.max_length,
+                    use_cache=True,
+                    **kwargs
                     )
-            scores = torch.stack(dec_output.scores, 1)
-            cache.update({'scores': scores})
-            return dec_output.sequences, cache
-
-        decoder_input_ids = self._shift_right(batch['sentence2_input_ids'])
-        decoder_inputs_embeds = self.shared(decoder_input_ids)
-        decoder_att_mask = batch['sentence2_attention_mask']
-        labels = batch['sentence2_input_ids'].clone()
-        labels[labels == self.pad_token_id] = -100
-
-        if self.args.combine_method == 'decoder_concat':
-            if self.args.ling2_only:
-                decoder_inputs_embeds = torch.cat([sent2_ling, decoder_inputs_embeds], dim=1)
-                decoder_att_mask = torch.cat([torch.ones((bs, 1)).to(decoder_inputs_embeds.device), decoder_att_mask], dim=1)
-                labels = torch.cat([torch.ones((bs, 1), dtype=torch.int64).to(decoder_inputs_embeds.device) * self.pad_token_id,
-                    labels], dim=1)
-            else:
-                decoder_inputs_embeds = torch.cat([sent1_ling, sent2_ling, decoder_inputs_embeds], dim=1)
-                decoder_att_mask = torch.cat([torch.ones((bs, 2)).to(decoder_inputs_embeds.device), decoder_att_mask], dim=1)
-                labels = torch.cat([torch.ones((bs, 2), dtype=torch.int64).to(decoder_inputs_embeds.device) * self.pad_token_id,
-                    labels], dim=1)
-        elif self.args.combine_method == 'decoder_add' or self.args.combine_method == 'decoder_add_first' :
-            if self.args.ling2_only:
-                decoder_inputs_embeds = decoder_inputs_embeds + self.args.combine_weight * sent2_ling
-            else:
-                decoder_inputs_embeds = decoder_inputs_embeds + sent1_ling + sent2_ling
+            return dec_output, cache
 
-        dec_output = self(
+        if sentence2_input_ids is not None:
+            labels = sentence2_input_ids.clone()
+            labels[labels == self.pad_token_id] = -100
+        else:
+            labels = None
+
+        if decoder_inputs_embeds is None:
+            decoder_input_ids = self._shift_right(sentence2_input_ids)
+            decoder_inputs_embeds = self.shared(decoder_input_ids)
+
+            if self.args.combine_method == 'decoder_concat':
+                if self.args.ling2_only:
+                    decoder_inputs_embeds = torch.cat([sentence2_ling, decoder_inputs_embeds], dim=1)
+                    decoder_attention_mask = torch.cat([torch.ones((bs, 1)).to(decoder_inputs_embeds.device), decoder_attention_mask], dim=1)
+                    labels = torch.cat([torch.ones((bs, 1), dtype=torch.int64).to(decoder_inputs_embeds.device) * self.pad_token_id,
+                        labels], dim=1)
+                else:
+                    decoder_inputs_embeds = torch.cat([sentence1_ling, sentence2_ling, decoder_inputs_embeds], dim=1)
+                    decoder_attention_mask = torch.cat([torch.ones((bs, 2)).to(decoder_inputs_embeds.device), decoder_attention_mask], dim=1)
+                    labels = torch.cat([torch.ones((bs, 2), dtype=torch.int64).to(decoder_inputs_embeds.device) * self.pad_token_id,
+                        labels], dim=1)
+            elif self.args.combine_method == 'decoder_add' or self.args.combine_method == 'decoder_add_first' :
+                if self.args.ling2_only:
+                    decoder_inputs_embeds = decoder_inputs_embeds + self.args.combine_weight * sentence2_ling
+                else:
+                    decoder_inputs_embeds = decoder_inputs_embeds + sentence1_ling + sentence2_ling
+
+        if ling_embed is None:
+            ling_embed = sentence2_ling
+
+        dec_output = super().forward(
                 decoder_inputs_embeds=decoder_inputs_embeds,
-                decoder_attention_mask=decoder_att_mask,
-                encoder_outputs=enc_output,
-                attention_mask=inputs_att_mask,
+                decoder_attention_mask=decoder_attention_mask,
+                encoder_outputs=encoder_outputs,
+                attention_mask=encoder_attention_mask,
                 labels=labels,
+                ling_embed=ling_embed,
+                **kwargs
                 )
         if self.args.combine_method == 'logits_add':
-            dec_output.logits = dec_output.logits + self.args.combine_weight * sent2_ling
+            dec_output.logits = dec_output.logits + self.args.combine_weight * sentence2_ling
             vocab_size = dec_output.logits.size(-1)
             dec_output.loss = F.cross_entropy(dec_output.logits.view(-1, vocab_size), labels.view(-1))
         return dec_output, cache
 
+    def generate(self, *args, **kwargs):
+        return self.forward(*args, **kwargs, generate=True)
+
+
+    def forward(self,
+                input_ids=None,
+                attention_mask=None,
+                labels=None,
+                decoder_attention_mask=None,
+                decoder_inputs_embeds=None,
+                sentence1_ling=None,
+                sentence2_ling=None,
+                sentence1_ling_embed=None,
+                sentence2_ling_embed=None,
+                inputs_embeds=None,
+                generate=False,
+                encoder_outputs=None,
+                encoder_attention_mask=None,
+                ling_embed=None,
+                generate_with_grad=False,
+                **kwargs):
 
-    def convert(self, batch, generate=False):
-        enc_output, enc_att_mask, cache = self.encode(batch)
-        dec_output, cache2 = self.decode(batch, enc_output, enc_att_mask, generate)
+        cache = {}
+        if encoder_outputs is None:
+            encoder_outputs, encoder_attention_mask, cache = self.encode(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                sentence1_ling=sentence1_ling,
+                sentence2_ling=sentence2_ling,
+                sentence1_ling_embed=sentence1_ling_embed,
+                sentence2_ling_embed=sentence2_ling_embed,
+                inputs_embeds=inputs_embeds
+            )
+        
+        dec_output, cache2 = self.decode(
+            sentence2_input_ids=labels,
+            sentence1_ling=sentence1_ling,
+            sentence2_ling=sentence2_ling,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            decoder_attention_mask=decoder_attention_mask,
+            encoder_outputs=encoder_outputs,
+            encoder_attention_mask=encoder_attention_mask,
+            generate=generate,
+            sentence1_ling_embed=sentence1_ling_embed,
+            sentence2_ling_embed=sentence2_ling_embed,
+            ling_embed=ling_embed,
+            generate_with_grad=generate_with_grad,
+            **kwargs
+        )
+        
         cache.update(cache2)
-        return dec_output, enc_output, cache
+        if generate:
+            return dec_output
+        else:
+            return MySeq2SeqLMOutput(
+                loss=dec_output.loss,
+                logits=dec_output.logits,
+                past_key_values=dec_output.past_key_values,
+                decoder_hidden_states=dec_output.decoder_hidden_states,
+                decoder_attentions=dec_output.decoder_attentions,
+                cross_attentions=dec_output.cross_attentions,
+                encoder_last_hidden_state=encoder_outputs[0],
+                encoder_hidden_states=getattr(encoder_outputs, 'hidden_states', None),
+                encoder_attentions=getattr(encoder_outputs, 'attentions', None),
+                cache=cache
+                )
 
     def infer_with_cache(self, batch):
-        dec_output, _, cache = self.convert(batch, generate = True)
+        dec_output, _, cache = self(batch, generate = True)
         return dec_output, cache
 
     def infer(self, batch):
         dec_output, _ = self.infer_with_cache(batch)
         return dec_output
 
-    def infer_with_feedback_BP(self, ling_disc, sem_emb, batch, tokenizer):
+    def infer_with_feedback_BP(self, ling_disc, sem_emb, batch, tokenizer, progress=None):
         from torch.autograd import grad
         interpolations = []
         def line_search():
-            best_val = None
-            best_loss = None
             eta = 1e3
             sem_prob = 1
             patience = 4
@@ -516,13 +634,11 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
                     new_loss, pred = get_loss(param_)
                 max_len = pred.shape[1]
                 lens = torch.where(pred == self.eos_token_id, 1, 0).argmax(-1) + 1
-                batch.update({
-                    'sentence2_input_ids': pred,
-                    'sentence2_attention_mask': sequence_mask(lens, max_len = max_len)
-                    })
-                sem_prob = torch.sigmoid(sem_emb.compare_sem(**batch)).item()
-                # if sem_prob <= 0.1:
-                #     patience -= 1
+                sem_batch = {**batch,
+                             'sentence2_input_ids': pred,
+                             'sentence2_attention_mask': sequence_mask(lens, max_len = max_len)
+                             }
+                sem_prob = torch.sigmoid(sem_emb.compare_sem(**sem_batch)).item()
                 if new_loss < loss and sem_prob >= 0.90 and lens.item() > 1:
                     return param_
                 eta *= 2.25
@@ -531,7 +647,7 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
 
         def get_loss(param):
             if self.args.feedback_param == 'l':
-                batch.update({'sent2_ling_embed': param})
+                batch.update({'sentence2_ling_embed': param})
             elif self.args.feedback_param == 's':
                 batch.update({'inputs_embeds': param})
 
@@ -539,8 +655,9 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
                 logits = param
                 pred = param.argmax(-1)
             else:
-                pred, cache = self.infer_with_cache(batch)
-                logits = cache['scores']
+                outputs = self.generate(**batch, output_scores=True, return_dict_in_generate=True, generate_with_grad=True)
+                pred = outputs.sequences
+                logits = torch.stack(outputs.scores, dim=1)
             out = ling_disc(logits = logits)
             probs = F.softmax(out, 1)
             if ling_disc.quant:
@@ -553,13 +670,13 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
             ling2_embed = self.ling_embed(batch['sentence2_ling'])
             param = torch.nn.Parameter(ling2_embed, requires_grad = True)
         elif self.args.feedback_param == 's':
-            inputs_embeds = self.shared(batch['sentence1_input_ids'])
+            inputs_embeds = self.shared(batch['input_ids'])
             param = torch.nn.Parameter(inputs_embeds, requires_grad = True)
         elif self.args.feedback_param == 'logits':
             logits = self.infer_with_cache(batch)[1]['scores']
             param = torch.nn.Parameter(logits, requires_grad = True)
-        target_np = batch['sentence2_ling'][0].cpu().numpy()
-        while True:
+        num_iter = 0
+        while num_iter < 3:
             loss, pred = get_loss(param)
             pred_text = tokenizer.batch_decode(pred.cpu().numpy(),
                     skip_special_tokens=True)[0]
@@ -571,6 +688,9 @@ class EncoderDecoderVAE(T5ForConditionalGeneration):
             param = line_search()
             if param is False:
                 break
+            num_iter += 1
+            if progress is not None:
+                progress((num_iter, None), unit='intermediate paraphrase generated.')
         return pred, [pred_text, interpolations]
 
 def set_grad(module, state):
@@ -609,7 +729,7 @@ class LingDiscPipeline():
     def __init__(self,
                  model_name="google/flan-t5-base",
                  disc_type='deberta',
-                 disc_ckpt='/data/mohamed/checkpoints/ling_disc/deberta-v3-small_flan-t5-base_40',
+                 disc_ckpt='mohdelgaar/lingconv-discriminator',
                  # disc_type='t5',
                  # disc_ckpt='/data/mohamed/checkpoints/ling_conversion_ling_disc.pt',
                  ):
@@ -629,15 +749,13 @@ def get_model(args, tokenizer, device):
         ling_disc = LingDisc(args.model_name, args.disc_type, args.disc_model_path).to(device)
     else:
         ling_disc = None
-    if args.linggen_type != 'none':
-        ling_gen = LingGenerator(args).to(device)
 
-    if not args.pretrain_disc:
+    if args.model_path:
         model = EncoderDecoderVAE.from_pretrained(args.model_path, args, tokenizer.pad_token_id, tokenizer.eos_token_id).to(device)
     else:
-        model = ling_disc
+        model = EncoderDecoderVAE.from_pretrained(args.model_name, args, tokenizer.pad_token_id, tokenizer.eos_token_id).to(device)
 
-    if args.sem_loss or args.sem_ckpt:
+    if args.sem_loss or args.model_path:
         if args.sem_loss_type == 'shared':
             sem_emb = model.encoder
         elif args.sem_loss_type == 'dedicated':
@@ -649,3 +767,14 @@ def get_model(args, tokenizer, device):
 
     return model, ling_disc, sem_emb
 
+@dataclass
+class MySeq2SeqLMOutput(Seq2SeqLMOutput):
+    """
+    Extends Seq2SeqLMOutput to include a cache dictionary for additional model outputs.
+    
+    Args:
+        cache (`Dict[str, Any]`):
+            Dictionary containing additional model outputs like linguistic features,
+            VAE parameters, scores, etc.
+    """
+    cache: Optional[Dict[str, Any]] = None

options.py

@@ -1,16 +1,28 @@
-import os, json
 import argparse
-import numpy as np
 from datetime import datetime
 from const import lftkplus_names
+import os, json
 from copy import deepcopy
+import numpy as np
 
+def str2bool(v):
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
 
 def parse_args(ckpt=None):
     parser = argparse.ArgumentParser()
+    parser.add_argument('--do_train', action='store_true')
+    parser.add_argument('--do_eval', action='store_true')
+    parser.add_argument('--do_predict', action='store_true')
     parser.add_argument('--data_dir', default='/data/mohamed/data')
     parser.add_argument('--data', default='ling_conversion')
-    parser.add_argument('--data_sources')
+    parser.add_argument('--data_sources', default='qqp,mrpc,stsb')
     parser.add_argument('--data_type', default='text')
     parser.add_argument('--aim_repo', default='/data/mohamed/')
     parser.add_argument('--ckpt_dir', default='/data/mohamed/checkpoints')
@@ -25,7 +37,7 @@ def parse_args(ckpt=None):
     parser.add_argument('--sem_loss_tao', default=0.5, type=float)
     parser.add_argument('--sem_loss_eps', default=1, type=float)
     parser.add_argument('--ckpt')
-    parser.add_argument('--disc_ckpt')
+    parser.add_argument('--disc_ckpt', default='mohdelgaar/lingconv-discriminator')
     parser.add_argument('--sem_ckpt')
     parser.add_argument('--lng_ids')
     parser.add_argument('--lng_ids_idx', type=int)
@@ -36,30 +48,34 @@ def parse_args(ckpt=None):
     parser.add_argument('--sem_path', default="mohdelgaar/lingconv-semantic-classifier")
     parser.add_argument('--sem_model_path', default="mohdelgaar/lingconv-semantic-classifier")
     parser.add_argument('--disc_model_path', default="mohdelgaar/lingconv-discriminator")
-    parser.add_argument('--disc_type', default="t5")
-    parser.add_argument('--aim_exp', default='ling-conversion')
+    parser.add_argument('--disc_type', default="deberta")
+    parser.add_argument('--aim_exp', default='lingconv-1201')
     parser.add_argument('--sem_loss_type', default='dedicated')
-    parser.add_argument('--combine_method', default='none')
+    parser.add_argument('--combine_method', default='decoder_add_first')
+    parser.add_argument('--injection_type', default='first')
+    parser.add_argument('--injection_layer', type=int, default=1)
     parser.add_argument('--train_log', type=int, default=200)
-    parser.add_argument('--val_log', type=int, default=2000)
-    parser.add_argument('--batch_size', type=int, default=64)
-    parser.add_argument('--eval_batch_size', type=int, default=32)
-    parser.add_argument('--max_eval_samples', type=int, default=1000)
-    parser.add_argument('--test_batch_size', type=int, default=1)
+    parser.add_argument('--val_log', type=int, default=1000)
+    parser.add_argument('--warmup_steps', type=int, default=1000)
+    parser.add_argument('--batch_size', type=int, default=16)
+    parser.add_argument('--eval_batch_size', type=int, default=256)
+    parser.add_argument('--max_eval_samples', type=int, default=3000)
+    parser.add_argument('--test_batch_size', type=int, default=256)
     parser.add_argument('--hidden_dim', type=int, default=500)
     parser.add_argument('--latent_dim', type=int, default=150)
     parser.add_argument('--lng_dim', type=int, default=40)
-    parser.add_argument('--disc_lng_dim', type=int)
+    parser.add_argument('--disc_lng_dim', type=int, default=40)
     parser.add_argument('--use_lora', action='store_true')
     parser.add_argument('--lora_r', type=int, default=64)
     parser.add_argument('--gpu', type=str, default='0')
-    parser.add_argument('--epochs', type=int, default=10)
+    parser.add_argument('--epochs', type=int, default=2)
     parser.add_argument('--grad_accumulation', type=int, default=1)
     parser.add_argument('--n_ica', type=int, default=10)
     parser.add_argument('--max_length', type=int, default=200)
     parser.add_argument('--total_steps', type=int)
     parser.add_argument('--kld_const', type=float, default=1)
-    parser.add_argument('--lr', type=float, default=1e-4)
+    parser.add_argument('--lr', type=float, default=1e-3)
+    parser.add_argument('--initializer_range', type=float, default=0.02)
     parser.add_argument('--kl_weight', type=float, default=1e-1)
     parser.add_argument('--weight_decay', type=float, default=1e-2)
     parser.add_argument('--ling_dropout', type=float, default=0.1)
@@ -71,12 +87,12 @@ def parse_args(ckpt=None):
     parser.add_argument('--pretrain_disc', action='store_true')
     parser.add_argument('--linggen_type', default='none')
     parser.add_argument('--linggen_input', default='s+l')
-    parser.add_argument('--aug_same', action='store_true')
+    parser.add_argument("--aug_same", type=str2bool, nargs='?', const=True, default=False)
     parser.add_argument('--ling_vae', action='store_true')
     parser.add_argument('--process_lingpred', action='store_true')
     parser.add_argument('--fudge_lambda', type=float, default=1.0)
     parser.add_argument('--use_lingpred', action='store_true')
-    parser.add_argument('--ling2_only', action='store_true')
+    parser.add_argument('--ling2_only', action='store_true', default=True)
     parser.add_argument('--cycle_loss', action='store_true')
     parser.add_argument('--disc_loss', action='store_true')
     parser.add_argument('--sem_loss', action='store_true')
@@ -96,19 +112,36 @@ def parse_args(ckpt=None):
     parser.add_argument('--quant_nbins', type=int, default=20)
     parser.add_argument('--src_lng', default = 'ling')
     parser.add_argument('--to_restore', nargs='+', default=[])
+    parser.add_argument('--freeze_lm', action='store_true', 
+                       help='Freeze the language model and only train the linguistic embedding')
+    parser.add_argument('--prepend_prompt', action='store_true',
+                       help='Prepend "generate a paraphrase: " to input text')
+    parser.add_argument('--prompt_text', type=str, default="generate a paraphrase: ",
+                       help='Text to prepend to input if prepend_prompt is True')
+    parser.add_argument('--do_imputation', action='store_true',
+                       help='Whether to perform imputation on linguistic features')
+    parser.add_argument('--imputation_percentage', type=int, default=20,
+                       help='Percentage of features to impute (20, 40, 60, 80)')
+    parser.add_argument('--imputation_seed', type=int, default=0,
+                       help='Seed for imputation set selection (0, 1, 2)')
     # args = parser.parse_args()
     args, unknown = parser.parse_known_args()
     args.name = f'{datetime.now().strftime("%m%d_%H-%M-%S")}-{args.data}-{args.combine_method}'
 
     major_arg = args.major_arg
     to_restore = [
+            'total_steps','major_arg','gpu','demo', 'eval_only', 'save_predict', 'predict_fn', 'fudge', 'predict_with_feedback',
+            'feedback_param', 'fb_log', 'data_dir', 'data', 'disc_ckpt', 'disc_type', 'sem_ckpt', 'fudge_lambda', 'eval_batch_size', 'test_batch_size', 'max_eval_samples',
+            'do_train', 'do_eval', 'do_predict', 
             ] + args.to_restore
     to_restore = {k: args.__dict__[k] for k in to_restore}
 
     if not args.disc_loss or args.disc_ckpt:
         args.disc_steps = 0
 
-    if args.data_sources is not None:
+    if args.data_sources == 'all':
+        args.data_sources = None
+    elif args.data_sources is not None:
         args.data_sources = args.data_sources.split(',')
 
     if ckpt is not None:
@@ -120,13 +153,17 @@ def parse_args(ckpt=None):
             ckpts = args.ckpt.split(',')
             args_list = [deepcopy(args) for _ in range(len(ckpts))]
             for i in range(len(ckpts)):
-                args_path = ckpts[i].replace('_best', '').replace('.pt', '.json')
+                args_path = ckpts[i].replace('_best', '').replace('.pt', '') + '.json'
                 with open(args_path) as f:
                     args_list[i].__dict__.update(json.load(f))
                 args_list[i].__dict__.update(to_restore)
                 args_list[i].ckpt = ckpts[i]
         else:
-            args_path = args.ckpt.replace('_best', '').replace('.pt', '.json')
+            args.ckpt = args.ckpt.rstrip('/')
+            if 'checkpoint-' in args.ckpt:
+                args_path = os.path.dirname(args.ckpt) + '.json'
+            else:
+                args_path = args.ckpt.replace('.pt', '') + '.json'
             ckpt = args.ckpt
             with open(args_path) as f:
                 args.__dict__.update(json.load(f))