検索・分類を強化する「ModernBERT」長文対応可能なAIモデル徹底解説

!pip install triton
!pip install flash-attn
!pip install transformers
!pip install torch
!pip install git+https://github.com/huggingface/transformers.git

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

# デバイスの確認
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# モデルの読み込み（bfloat16を使用）
model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# テキストの処理
text = "The capital of France is [MASK]."
inputs = tokenizer(text, return_tensors="pt").to(device)

# 推論の実行（自動混合精度を使用）
with torch.cuda.amp.autocast():
    outputs = model(**inputs)
    
# マスクされた位置の予測を取得
mask_token_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
predicted_token_id = outputs.logits[0, mask_token_index].argmax(axis=-1)
predicted_token = tokenizer.decode(predicted_token_id)

print(f"Predicted word: {predicted_token}")

Predicted word:  Paris

import torch
from transformers import AutoTokenizer, AutoModel

model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")

def get_text_embedding(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True).to("cuda")
    with torch.amp.autocast('cuda'):
        outputs = model(**inputs)
        return outputs.last_hidden_state.mean(dim=1)

# 文書データベース
documents = [
    "Python is a popular programming language",
    "JavaScript is used for web development",
    "Machine learning uses statistical methods"
]

# 検索クエリ
query = "What programming languages are popular?"
query_emb = get_text_embedding(query)

# 類似度に基づいて関連文書を検索
docs_emb = torch.cat([get_text_embedding(doc) for doc in documents])
similarities = torch.cosine_similarity(query_emb, docs_emb)

# 最も関連性の高い文書を表示
top_idx = similarities.argmax().item()
print(f"Most relevant document: {documents[top_idx]}")

Asking to truncate to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no truncation.
Most relevant document: Python is a popular programming language

import torch
from transformers import AutoTokenizer, AutoModel

model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")

def search_code(query, code_snippets):
    query_emb = get_text_embedding(query)
    code_embs = torch.cat([get_text_embedding(code) for code in code_snippets])
    
    similarities = torch.cosine_similarity(query_emb, code_embs)
    top_idx = similarities.argmax().item()
    return code_snippets[top_idx]

# 使用例
code_snippets = [
    "def fibonacci(n):\n    return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)",
    "def quick_sort(arr):\n    if len(arr) <= 1: return arr",
    "def binary_search(arr, x):\n    left, right = 0, len(arr)-1"
]

query = "How to implement fibonacci sequence"
relevant_code = search_code(query, code_snippets)
print(f"Relevant code:\n{relevant_code}")

Relevant code:
def fibonacci(n):
    return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)

import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

# デバイスの確認
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# モデルの読み込み
model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

# マスクトークンを使用した入力テキスト
text = "The Eiffel Tower is located in [MASK]."
inputs = tokenizer(text, return_tensors="pt").to(device)

# 推論の実行
with torch.cuda.amp.autocast():
    outputs = model(**inputs)

# マスクされた位置の予測を取得
mask_token_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
predicted_token_id = outputs.logits[0, mask_token_index].argmax(axis=-1)
predicted_token = tokenizer.decode(predicted_token_id)

# 上位5つの予測を表示
topk = torch.topk(outputs.logits[0, mask_token_index], k=5)
top_tokens = [tokenizer.decode(token_id) for token_id in topk.indices]

print(f"入力文: {text}")
print(f"最も確率の高い予測: {predicted_token}")
print(f"上位5つの予測: {', '.join(top_tokens)}")

入力文: The Eiffel Tower is located in [MASK].
最も確率の高い予測:  Paris
上位5つの予測:  Paris,  France,  Switzerland,  Belgium,  Geneva

from transformers import BertTokenizer, BertForMaskedLM
import torch

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForMaskedLM.from_pretrained('bert-base-uncased')

# GPUが利用可能な場合は使用
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

def predict_masked_word(text):
    # 入力のトークン化
    inputs = tokenizer(text, return_tensors="pt")
    inputs = {k: v.to(device) for k, v in inputs.items()}

    # 予測の実行
    with torch.no_grad():
        outputs = model(**inputs)
    
    # マスクトークンの位置を特定
    mask_token_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
    
    # 上位5個の予測を取得
    top5_tokens = torch.topk(outputs.logits[0, mask_token_index], 5)
    predictions = [tokenizer.decode([token_id]) for token_id in top5_tokens.indices]
    
    return predictions

# テスト
text = "The Eiffel Tower is located in [MASK]."
predictions = predict_masked_word(text)
print(f"Top 5 predictions: {predictions}")

Top 5 predictions: ['paris', 'luxembourg', 'brussels', 'france', 'monaco']

import torch
from transformers import AutoTokenizer, AutoModel
from torch.nn.functional import cosine_similarity

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# モデルの読み込み - AutoModel を使用
model_id = "answerdotai/ModernBERT-base"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModel.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

def get_embeddings(sentences):
    # バッチ処理で効率化
    inputs = tokenizer(
        sentences, 
        padding=True, 
        truncation=True, 
        max_length=512, 
        return_tensors="pt"
    ).to(device)
    
    with torch.amp.autocast('cuda'):
        outputs = model(**inputs)
        # 平均プーリングで文の埋め込みを取得
        embeddings = outputs.last_hidden_state.mean(dim=1)
    return embeddings

# テスト文
test_pairs = [
    ("Paris is the capital of France.", "The main city of France is Paris."),
    ("I love eating apples.", "Apples are my favorite fruit."),
    ("The cat is sleeping.", "The dog is running."),
]

# 各ペアの類似度を計算
for sent1, sent2 in test_pairs:
    # 両方の文の埋め込みを一度に取得
    embeddings = get_embeddings([sent1, sent2])
    emb1, emb2 = embeddings[0], embeddings[1]
    
    similarity = cosine_similarity(
        emb1.unsqueeze(0), 
        emb2.unsqueeze(0)
    ).item()
    
    print(f"\nSentence 1: {sent1}")
    print(f"Sentence 2: {sent2}")
    print(f"Similarity: {similarity:.4f}")

Sentence 1: Paris is the capital of France.
Sentence 2: The main city of France is Paris.
Similarity: 0.9855

Sentence 1: I love eating apples.
Sentence 2: Apples are my favorite fruit.
Similarity: 0.9503

Sentence 1: The cat is sleeping.
Sentence 2: The dog is running.
Similarity: 0.9826

from transformers import BertTokenizer, BertModel
import torch
from torch.nn.functional import cosine_similarity

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model = model.to(device)

def get_sentence_embedding(text):
    # トークン化
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    inputs = {k: v.to(device) for k, v in inputs.items()}
    
    # 埋め込みの取得
    with torch.no_grad():
        outputs = model(**inputs)
    
    # [CLS]トークンの出力を使用
    return outputs.last_hidden_state[:, 0, :]

# テスト文
sent1 = "I love eating apples."
sent2 = "Apples are my favorite fruit."

# 埋め込みの取得と類似度の計算
emb1 = get_sentence_embedding(sent1)
emb2 = get_sentence_embedding(sent2)
similarity = cosine_similarity(emb1, emb2)

print(f"Similarity score: {similarity.item():.4f}")

Similarity score: 0.9593