""" TransformersAI - Hugging Face Transformersモデル用アダプター Llama 3.2、Qwen、Mistral、Gemma等のローカルモデルに対応 """ from typing import List, Tuple, Any, Optional import os from .base import BaseAI class TransformersAI(BaseAI): """ Hugging Face Transformersモデル用アダプター 特徴: - ローカルでモデルをロード - logitsから直接確率を取得可能 - user/assistantを明確に分離する形式を推奨（Llama 3.2形式） """ _instances = {} # モデルパスごとのインスタンスをキャッシュ（常駐） def __new__(cls, model_path: str = None): """シングルトンパターンでモデルを常駐""" path = model_path or os.getenv("HF_MODEL_REPO", "meta-llama/Llama-3.2-3B-Instruct") if path not in cls._instances: cls._instances[path] = super().__new__(cls) cls._instances[path]._initialized = False return cls._instances[path] def __init__(self, model_path: str = None): """ モデルをロードして初期化（一度だけ実行、常駐） Args: model_path: モデルリポジトリIDまたはローカルパス """ if hasattr(self, '_initialized') and self._initialized: return self.model_path = model_path or os.getenv("HF_MODEL_REPO", "meta-llama/Llama-3.2-3B-Instruct") self.model = self._load_model(self.model_path) self._initialized = True if self.model is None: raise ValueError(f"モデルのロードに失敗しました: {self.model_path}") @classmethod def get_model(cls, model_path: str = None) -> 'TransformersAI': """モデルインスタンスを取得（常駐キャッシュから）""" return cls(model_path) @classmethod def clear_cache(cls): """キャッシュをクリア（開発・テスト用）""" cls._instances.clear() def _load_model(self, model_path: str) -> Optional[Any]: """モデルをロード（Transformers使用、Hubから直接読み込み）""" try: if not model_path: return None # モデルパスがリポジトリID（"user/repo"形式）か、ローカルパスかを判定 is_repo_id = "/" in model_path and not os.path.exists(model_path) # リポジトリIDの場合は os.path.exists() チェックをスキップ if not is_repo_id and not os.path.exists(model_path): print(f"[TransformersAI] モデルパスが存在しません: {model_path}") return None # transformersを使用してモデルをロード try: from transformers import AutoModelForCausalLM, AutoTokenizer import torch # GPUが利用可能かチェック device = "cuda" if torch.cuda.is_available() else "cpu" if device == "cuda": print("[TransformersAI] GPU検出: CUDAを使用します") else: print("[TransformersAI] GPU未検出: CPUモードで実行します") print(f"[TransformersAI] モデルをロード中: {model_path}") print(f"[TransformersAI] デバイス: {device}") hf_token = os.getenv("HF_TOKEN") if is_repo_id: print(f"[TransformersAI] Hugging Face Hub から直接読み込み: {model_path}") else: print(f"[TransformersAI] ローカルパスから読み込み: {model_path}") # トークナイザーとモデルをロード（Hubから直接読み込む） tokenizer = AutoTokenizer.from_pretrained( model_path, token=hf_token, ) # device_map="auto"を使用する場合はaccelerateが必要 # accelerateがインストールされていない場合は、device_mapをNoneにして手動でデバイスに移動 try: import accelerate use_device_map = "auto" if device == "cuda" else None except ImportError: print("[TransformersAI] 警告: accelerateがインストールされていません。device_mapを使用しません。") use_device_map = None model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.float16 if device == "cuda" else torch.float32, device_map=use_device_map, token=hf_token, ) # device_mapを使用しない場合は、手動でデバイスに移動 if use_device_map is None: model = model.to(device) # モデルとトークナイザーをタプルで返す print(f"[TransformersAI] モデルロード成功 ({device}モード)") return (model, tokenizer) except Exception as e: import traceback print(f"[TransformersAI] transformersでのロードに失敗: {e}") traceback.print_exc() return None except Exception as e: import traceback print(f"[TransformersAI] モデルロードエラー: {e}") traceback.print_exc() return None def get_token_probabilities(self, text: str, k: int = 5) -> List[Tuple[str, float]]: """ 文章とkを引数に、{token, 確率}のリストを返す Args: text: 入力文章 k: 取得するトークン数 Returns: List[Tuple[str, float]]: (トークン, 確率)のリスト """ if self.model is None: return [] try: # transformers モデルの場合 if isinstance(self.model, tuple) and len(self.model) == 2: model, tokenizer = self.model import torch # テキストをトークン化 inputs = tokenizer(text, return_tensors="pt") device = next(model.parameters()).device inputs = {k: v.to(device) for k, v in inputs.items()} # モデルで推論（勾配計算なし） with torch.no_grad(): outputs = model(**inputs) logits = outputs.logits[0, -1, :] # 最後のトークンのlogits # logitsを確率に変換（softmax） probs = torch.softmax(logits, dim=-1) # 上位k個のトークンを取得 top_probs, top_indices = torch.topk(probs, k) # トークンIDを文字列に変換 items: List[Tuple[str, float]] = [] # 特殊トークンを定義（Llama 3.2、Qwen、Mistral等で使用） SPECIAL_TOKENS = [ "<|begin_of_text|>", "<|end_of_text|>", "<|eot_id|>", "<|start_header_id|>", "<|end_header_id|>", "<|im_start|>", "<|im_end|>", ] def _clean_text_local(text: str) -> str: """制御文字・不可視文字・置換文字・特殊トークンを厳密に取り除く""" if not text: return "" # 特殊トークンを除去 for special_token in SPECIAL_TOKENS: text = text.replace(special_token, "") # 基底クラスの_clean_textを使用 return self._clean_text(text) for idx, prob in zip(top_indices, top_probs): token_id = idx.item() # skip_special_tokens=Trueで特殊トークンを除外 token = tokenizer.decode([token_id], skip_special_tokens=True, clean_up_tokenization_spaces=False) token = _clean_text_local(token) # 空文字列のトークンは除外 if not token: continue prob_value = prob.item() items.append((token, float(prob_value))) # 確率を正規化 if items: total_prob = sum(prob for _, prob in items) if total_prob > 0: normalized_items: List[Tuple[str, float]] = [] for token, prob in items: normalized_prob = prob / total_prob normalized_items.append((token, normalized_prob)) return normalized_items return items else: print("[TransformersAI] モデルがサポートされていません") return [] except Exception as e: print(f"[TransformersAI] トークン確率取得エラー: {e}") import traceback traceback.print_exc() return [] def build_chat_prompt( self, user_content: str, system_content: str = "あなたは親切で役に立つAIアシスタントです。簡潔な回答をしてください", assistant_content: Optional[str] = None ) -> str: """ チャットプロンプトを構築（Llama 3.2形式） 注意: Transformersモデル（特にLlama 3.2、Qwen等）では、 user/assistantを明確に分離する形式を推奨します。 Args: user_content: ユーザーのメッセージ system_content: システムプロンプト assistant_content: アシスタントの既存応答（会話履歴用、オプション） Returns: str: Llama 3.2形式のプロンプト """ # 既に整形済みのプロンプトが渡されている場合（複数行、ヘッダーを含む） # そのまま返す if "<|start_header_id|>" in user_content or "<|eot_id|>" in user_content: return user_content # Llama 3.2形式でプロンプトを構築 prompt_parts = [] # Systemメッセージ if system_content: prompt_parts.append("<|start_header_id|>system<|end_header_id|>") prompt_parts.append(system_content) prompt_parts.append("<|eot_id|>") # Userメッセージ prompt_parts.append("<|start_header_id|>user<|end_header_id|>") prompt_parts.append(user_content) prompt_parts.append("<|eot_id|>") # Assistantメッセージ（会話履歴がある場合） if assistant_content: prompt_parts.append("<|start_header_id|>assistant<|end_header_id|>") prompt_parts.append(assistant_content) prompt_parts.append("<|eot_id|>") # 新しい応答を生成する場合は、assistantヘッダーだけを追加 prompt_parts.append("<|start_header_id|>assistant<|end_header_id|>") prompt_text = "\n".join(prompt_parts) # BOS(<|begin_of_text|>) の重複を抑止: 先頭のBOSを全て除去 # transformers側でBOSが自動付与される場合があるため BOS = "<|begin_of_text|>" s = prompt_text.lstrip() while s.startswith(BOS): s = s[len(BOS):] prompt_text = s return prompt_text