1.はじめに

multilingual-e5-largeが埋め込み表現作成手段として優秀とのことだったので、使ってみたくなり、huggingfaceの公式のコードから出発。

import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel


def average_pool(last_hidden_states: Tensor,
                 attention_mask: Tensor) -> Tensor:
    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]


# Each input text should start with "query: " or "passage: ", even for non-English texts.
# For tasks other than retrieval, you can simply use the "query: " prefix.
input_texts = ['query: how much protein should a female eat',
               'query: 南瓜的家常做法',
               "passage: As a general guideline, the CDC's average requirement of protein for women ages 19 to 70 is 46 grams per day. But, as you can see from this chart, you'll need to increase that if you're expecting or training for a marathon. Check out the chart below to see how much protein you should be eating each day.",
               "passage: 1.清炒南瓜丝 原料:嫩南瓜半个 调料:葱、盐、白糖、鸡精 做法: 1、南瓜用刀薄薄的削去表面一层皮,用勺子刮去瓤 2、擦成细丝(没有擦菜板就用刀慢慢切成细丝) 3、锅烧热放油,入葱花煸出香味 4、入南瓜丝快速翻炒一分钟左右,放盐、一点白糖和鸡精调味出锅 2.香葱炒南瓜 原料:南瓜1只 调料:香葱、蒜末、橄榄油、盐 做法: 1、将南瓜去皮,切成片 2、油锅8成热后,将蒜末放入爆香 3、爆香后,将南瓜片放入,翻炒 4、在翻炒的同时,可以不时地往锅里加水,但不要太多 5、放入盐,炒匀 6、南瓜差不多软和绵了之后,就可以关火 7、撒入香葱,即可出锅"]

tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large')
model = AutoModel.from_pretrained('intfloat/multilingual-e5-large')

# Tokenize the input texts
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt')

outputs = model(**batch_dict)
embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

# normalize embeddings
embeddings = F.normalize(embeddings, p=2, dim=1)
scores = (embeddings[:2] @ embeddings[2:].T) * 100
print(scores.tolist())

が、公式のサイトの方法をそのままcolaboratoryで使うと
（１）GPUが使われてない(T4に設定変更設定しただけの状態)
（２）３〜４回実行するとGPUがいっぱいになって実行できなくなる。
こんなエラー

OutOfMemoryError: CUDA out of memory. Tried to allocate 166.00 MiB. GPU 0 has a total capacty of 14.75 GiB of which 51.06 MiB is free. Process 215610 has 14.70 GiB memory in use. Of the allocated memory 14.06 GiB is allocated by PyTorch, and 528.19 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

　という問題が。

　そこで上記２つの問題を解決するためにちょっと修正した。

2.変更点

（１）対策：GPUを使うための設定を追加

import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"

model = AutoModel.from_pretrained('intfloat/multilingual-e5-large').to(device)
・・・
batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt').to(device)

のように、modelとbatch_dictに.to(device)を追加。

上記２つに設定しないと、CPUとGPUにデータが分散しているのでダメ、と怒られる。

（２）GPUメモリを残さない（？）設定に変更
　ネット上を彷徨うとkillしなさいとかあるが、killするとcolaboratoryの実行環境自体が再起動してしまった。結局、@torch.no_grad()というデコレータをembedの関数に追加したところ、解消した。

最終的にこんな感じに。

import torch.nn.functional as F

from torch import Tensor
from transformers import AutoTokenizer, AutoModel


def average_pool(last_hidden_states: Tensor,
                attention_mask: Tensor) -> Tensor:
   last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
   return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

import torch
device = "cuda:0" if torch.cuda.is_available() else "cpu"

tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-large') 
model = AutoModel.from_pretrained('intfloat/multilingual-e5-large').to(device)


@torch.no_grad()
def e5large_embed(input_texts):
   # Tokenize the input texts  
   batch_dict = tokenizer(input_texts, max_length=512, padding=True, truncation=True, return_tensors='pt').to(device)

   outputs = model(**batch_dict)
   embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

   # normalize embeddings
   embeddings = F.normalize(embeddings, p=2, dim=1)

   del outputs
   del batch_dict
   torch.cuda.empty_cache()
   return embeddings
   #scores = (embeddings[:2] @ embeddings[2:].T) * 100
   #print(scores.tolist())

関数は

input_texts = frame["abst"].tolist()
embdata = e5large_embed(input_texts)

のような感じで使う。

これで何回実行しても、GPUのmemoryがいっぱいになってしまうことがなくなった。

3.其その他

同じと思われる書き方があった！こちらの方がいいかも。
https://www.pinecone.io/models/multilingual-e5-large/

これも注意

Do I need to add the prefix "query: " and "passage: " to input texts?
Yes, this is how the model is trained, otherwise you will see a performance degradation.

Here are some rules of thumb:

Use "query: " and "passage: " correspondingly for asymmetric tasks such as passage retrieval in open QA, ad-hoc information retrieval.

Use "query: " prefix for symmetric tasks such as semantic similarity, bitext mining, paraphrase retrieval.

Use "query: " prefix if you want to use embeddings as features, such as linear probing classification, clustering.

multilingual-e5-largeを使った埋め込み表現作成時の注意

1.はじめに

2.変更点

3.其その他