|
|
""" |
|
|
Mistral モデル実装 |
|
|
|
|
|
Mistral AI社のモデル |
|
|
Sliding Window Attention、GQA採用 |
|
|
""" |
|
|
from typing import List, Tuple |
|
|
|
|
|
import torch |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
from .base import BaseLanguageModel, ModelConfig |
|
|
|
|
|
|
|
|
|
|
|
MISTRAL_7B_CONFIG = ModelConfig( |
|
|
name="Mistral 7B v0.3", |
|
|
model_id="mistralai/Mistral-7B-v0.3", |
|
|
embedding_dim=4096, |
|
|
vocab_size=32768, |
|
|
) |
|
|
|
|
|
|
|
|
class MistralModel(BaseLanguageModel): |
|
|
""" |
|
|
Mistralモデルの実装 |
|
|
|
|
|
Mistral AI社のMistral 7Bシリーズ。 |
|
|
Sliding Window Attention、GQA採用。 |
|
|
""" |
|
|
|
|
|
|
|
|
LOGITS_NOISE_SCALE = 10.0 |
|
|
|
|
|
def load(self) -> None: |
|
|
"""モデルとトークナイザーをロード""" |
|
|
if self._is_loaded: |
|
|
return |
|
|
|
|
|
try: |
|
|
self._model = AutoModelForCausalLM.from_pretrained( |
|
|
self._config.model_id, |
|
|
torch_dtype="auto", |
|
|
) |
|
|
self._tokenizer = AutoTokenizer.from_pretrained(self._config.model_id) |
|
|
self._model.eval() |
|
|
self._is_loaded = True |
|
|
except Exception as e: |
|
|
raise RuntimeError(f"Failed to load model {self._config.model_id}: {e}") |
|
|
|
|
|
def forward_with_noise( |
|
|
self, noise: torch.Tensor |
|
|
) -> Tuple[torch.Tensor, torch.Tensor]: |
|
|
""" |
|
|
ノイズを入力として順伝播を実行し、出力にもノイズを加算 |
|
|
|
|
|
Args: |
|
|
noise: 入力ノイズテンソル [batch, seq_len, embedding_dim] |
|
|
|
|
|
Returns: |
|
|
Tuple[logits, corrupted_logits] |
|
|
""" |
|
|
if not self._is_loaded: |
|
|
raise RuntimeError("Model not loaded. Call load() first.") |
|
|
|
|
|
with torch.no_grad(): |
|
|
outputs = self._model(inputs_embeds=noise) |
|
|
logits = outputs.logits |
|
|
|
|
|
|
|
|
logits_noise = ( |
|
|
torch.randn_like(logits) * logits.std() * self.LOGITS_NOISE_SCALE |
|
|
) |
|
|
corrupted_logits = logits + logits_noise |
|
|
|
|
|
return logits, corrupted_logits |
|
|
|
|
|
def decode_indices(self, indices: List[int]) -> List[str]: |
|
|
"""トークンインデックスをデコード""" |
|
|
if not self._is_loaded: |
|
|
raise RuntimeError("Model not loaded. Call load() first.") |
|
|
|
|
|
return [self._tokenizer.decode([i]) for i in indices] |
|
|
|
