Upload COCOM

modelling_pisco.py

@@ -1036,11 +1036,12 @@ class COCOM(PreTrainedModel):
         """
         Compress a list of documents
         if questions is not None, assumes compression is done query-dependently !
+        excepts as many questions as documents here (so repeat questions for multidoc)
         """ 
         if questions is None:
             input_encoder = self.prepare_encoder_inputs(documents, max_length=128)
         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])
+            input_encoder = self.prepare_encoder_inputs(documents, max_length=128, q_texts=questions)
         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)