# BERT embeddings on big dataset

**URL:** </static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279%3E%5C
**Category:** 🤗Datasets\
**Created:** [August 27, 2024, 11:08pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279 "2024-08-27T23:08:58Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![simonberrebi](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/simonberrebi/32/30948_2.png) [@simonberrebi](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsimonberrebi)%5C
**Post date:** [August 27, 2024, 11:08pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279%2F1 "2024-08-27T23:08:59Z")

</div>

Hello friends, I am looking to get BERT embeddings on a dataset with ~20M rows. I’m able to get padded lists of tokens using:

```auto
model = BertModel.from_pretrained('bert-base-uncased')
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
tokenized_dataset = dataset_1.map(lambda x: tokenizer(x["charge_metadata__email"], padding = "longest"), batched=True)

```

But can’t get past this point. I am trying to turn the lists of tokens and attention masks into tensors for model inference but am getting this error:

`token_ids = tokenized_dataset.map(lambda x: torch.tensor(x["input_ids"]).unsqueeze(0), batched=True)`

> TypeError: Provided `function` which is applied to all elements of table returns a variable of type \<class ‘torch.Tensor’\>. Make sure provided `function` returns a variable of type `dict` (or a pyarrow table) to update the dataset or `None` if you are only interested in side effects.

Curious if this is the right approach.

---

<div class="post-metadata">

**Author:** ![mahmutc](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mahmutc/32/64626_2.png) [@mahmutc](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmahmutc)%5C
**Post date:** [August 28, 2024, 12:00pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279%2F2 "2024-08-28T12:00:46Z")

</div>

hi @simonberrebi  
I’m not sure if I follow well, but tokenized\_dataset is a `DatasetDict`, right?

If you want, you can remove irrelevant columns with: `tokenized_dataset.map(remove_columns=["blabla","blabla2"])`

But I’m not sure if you need it. If you’re trying to train a new model, why you don’t use `data_collator` as explained here::

> **[Fine-tuning a model with the Trainer API - Hugging Face NLP Course](https://huggingface.co/learn/nlp-course/chapter3/3)**
>
> We’re on a journey to advance and democratize artificial intelligence through open source and open science.

Edit:  
It looks you don’t even need data\_collator.

> Note that when you pass the `tokenizer` as we did here, the default `data_collator` used by the `Trainer` will be a `DataCollatorWithPadding` as defined previously, so you can skip the line `data_collator=data_collator` in this call.

---

<div class="post-metadata">

**Author:** ![mikehemberger](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/mikehemberger/32/28013_2.png) [@mikehemberger](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fmikehemberger)%5C
**Post date:** [August 28, 2024, 1:49pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279%2F3 "2024-08-28T13:49:19Z")

</div>

Hey @simonberrebi ,  
I think you should return a dict instead of a tensor (as the error msg suggests). So if you put your tensors inside a dict it might just work as expected. Let us know whether it worked 🙂  
Best,  
M

---

<div class="post-metadata">

**Author:** ![simonberrebi](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/simonberrebi/32/30948_2.png) [@simonberrebi](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fsimonberrebi)%5C
**Post date:** [August 28, 2024, 4:19pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Fbert-embeddings-on-big-dataset%2F104279%2F5 "2024-08-28T16:19:15Z")

</div>

That works great, than you friends!