import numpy as np
​
def preprocess_for_question_generation(examples):
    if examples["answer"] is None or examples["context"] is None:
        return None
    
    if not isinstance(examples["answer"], list) or not isinstance(examples["context"], list):
        return None
    
    if len(examples["answer"]) != len(examples["context"]):
        return None
    
    prepared_questions = []
    prepared_input_texts = []
    for context, answer_list in zip(examples["context"], examples["answer"]):
        for answer in answer_list:
            sentences = context.split(".")
            highlighted_sentences = []
            target_text = []
            for sentence in sentences:
                if answer.lower() in sentence.lower():
                    highlighted_sentence = f"<hl> {sentence.strip()} <hl>"
                    highlighted_sentences.append(highlighted_sentence)
                    target_text.extend(answer.split())
            if highlighted_sentences:
                input_text = " ".join(highlighted_sentences)
                prepared_input_texts.append(input_text)
                target_text = " <sep> ".join(target_text)
                prepared_questions.append(target_text)
    
    tokenized_inputs = tokenizer(
        prepared_input_texts,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    
    if "question" in examples: 
        tokenized_questions = tokenizer(
            prepared_questions,
            max_length=64,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        tokenized_inputs["labels"] = tokenized_questions["input_ids"].clone()
    
    # Ensure input_ids length is 128
    if len(tokenized_inputs["input_ids"][0]) != 128:
        tokenized_inputs["input_ids"] = torch.cat((tokenized_inputs["input_ids"], torch.zeros((1, 128 - len(tokenized_inputs["input_ids"][0])))), dim=1)
    
    # Convert tensors to numpy arrays
    tokenized_inputs_numpy = {}
    for key, value in tokenized_inputs.items():
        tokenized_inputs_numpy[key] = value.numpy()
    
    return tokenized_inputs_numpy
​
​
# Apply preprocessing to datasets
tokenized_train_dataset = train_dataset.map(preprocess_for_question_generation, batched=True)
tokenized_val_dataset = val_dataset.map(preprocess_for_question_generation, batched=True)
tokenized_test_dataset = test_dataset.map(preprocess_for_question_generation, batched=True)
0%
0/43 [00:00<?, ?ba/s]
---------------------------------------------------------------------------
ArrowInvalid                              Traceback (most recent call last)
Cell In[42], line 62
     58     return tokenized_inputs_numpy
     61 # Apply preprocessing to datasets
---> 62 tokenized_train_dataset = train_dataset.map(preprocess_for_question_generation, batched=True)
     63 tokenized_val_dataset = val_dataset.map(preprocess_for_question_generation, batched=True)
     64 tokenized_test_dataset = test_dataset.map(preprocess_for_question_generation, batched=True)

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:1955, in Dataset.map(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint, desc)
   1952 disable_tqdm = not logging.is_progress_bar_enabled()
   1954 if num_proc is None or num_proc == 1:
-> 1955     return self._map_single(
   1956         function=function,
   1957         with_indices=with_indices,
   1958         with_rank=with_rank,
   1959         input_columns=input_columns,
   1960         batched=batched,
   1961         batch_size=batch_size,
   1962         drop_last_batch=drop_last_batch,
   1963         remove_columns=remove_columns,
   1964         keep_in_memory=keep_in_memory,
   1965         load_from_cache_file=load_from_cache_file,
   1966         cache_file_name=cache_file_name,
   1967         writer_batch_size=writer_batch_size,
   1968         features=features,
   1969         disable_nullable=disable_nullable,
   1970         fn_kwargs=fn_kwargs,
   1971         new_fingerprint=new_fingerprint,
   1972         disable_tqdm=disable_tqdm,
   1973         desc=desc,
   1974     )
   1975 else:
   1977     def format_cache_file_name(cache_file_name, rank):

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:520, in transmit_tasks.<locals>.wrapper(*args, **kwargs)
    518     self: "Dataset" = kwargs.pop("self")
    519 # apply actual function
--> 520 out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    521 datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    522 for dataset in datasets:
    523     # Remove task templates if a column mapping of the template is no longer valid

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:487, in transmit_format.<locals>.wrapper(*args, **kwargs)
    480 self_format = {
    481     "type": self._format_type,
    482     "format_kwargs": self._format_kwargs,
    483     "columns": self._format_columns,
    484     "output_all_columns": self._output_all_columns,
    485 }
    486 # apply actual function
--> 487 out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    488 datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    489 # re-apply format to the output

File /opt/conda/lib/python3.10/site-packages/datasets/fingerprint.py:458, in fingerprint_transform.<locals>._fingerprint.<locals>.wrapper(*args, **kwargs)
    452             kwargs[fingerprint_name] = update_fingerprint(
    453                 self._fingerprint, transform, kwargs_for_fingerprint
    454             )
    456 # Call actual function
--> 458 out = func(self, *args, **kwargs)
    460 # Update fingerprint of in-place transforms + update in-place history of transforms
    462 if inplace:  # update after calling func so that the fingerprint doesn't change if the function fails

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:2356, in Dataset._map_single(self, function, with_indices, with_rank, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, disable_tqdm, desc, cache_only)
   2354                 writer.write_table(batch)
   2355             else:
-> 2356                 writer.write_batch(batch)
   2357 if update_data and writer is not None:
   2358     writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_writer.py:510, in ArrowWriter.write_batch(self, batch_examples, writer_batch_size)
    508     inferred_features[col] = typed_sequence.get_inferred_type()
    509 schema = inferred_features.arrow_schema if self.pa_writer is None else self.schema
--> 510 pa_table = pa.Table.from_arrays(arrays, schema=schema)
    511 self.write_table(pa_table, writer_batch_size)

File /opt/conda/lib/python3.10/site-packages/pyarrow/table.pxi:3674, in pyarrow.lib.Table.from_arrays()

File /opt/conda/lib/python3.10/site-packages/pyarrow/table.pxi:2837, in pyarrow.lib.Table.validate()

File /opt/conda/lib/python3.10/site-packages/pyarrow/error.pxi:100, in pyarrow.lib.check_status()

ArrowInvalid: Column 5 named input_ids expected length 1000 but got length 962