Upload COCOM

modelling_pisco.py

@@ -985,7 +985,7 @@ class COCOM(PreTrainedModel):
         # Creating encoder inputs:
         if query_dependent:
             # We provide the question for compression:
-            input_encoder = self.prepare_encoder_inputs(flat_documents, max_length=128, q_texts=questions)
+            input_encoder = self.prepare_encoder_inputs(flat_documents, max_length=128, q_texts=[question for question, docs in zip(questions, documents) for _ in docs])
         else:
             input_encoder = self.prepare_encoder_inputs(flat_documents, max_length=128)
             
@@ -1036,11 +1036,11 @@ class COCOM(PreTrainedModel):
         """
         Compress a list of documents
         if questions is not None, assumes compression is done query-dependently !
-        """
+        """ 
         if questions is None:
             input_encoder = self.prepare_encoder_inputs(documents, max_length=128)
-        else:
-            input_encoder = self.prepare_encoder_inputs(documents, max_length=128, q_texts=questions)
+        else: # we assume query-dependent here:
+            input_encoder = self.prepare_encoder_inputs(documents, max_length=128, q_texts=[question for question, docs in zip(questions, documents) for _ in docs])
         enc_input_ids = input_encoder['input_ids'].to(self.decoder.device)
         attention_mask = input_encoder['attention_mask'].to(self.decoder.device)
         return self.compress(enc_input_ids=enc_input_ids, enc_attention_mask=attention_mask)