from typing import List, Tuple, Any, Optional import os import subprocess from config import Config class AI: """AIクラス - モデルをロードして文章とkを引数にトークンと確率のリストを返す（常駐版）""" _instances = {} # モデルパスごとのインスタンスをキャッシュ（常駐） def __new__(cls, model_path: str = None): """シングルトンパターンでモデルを常駐""" path = model_path or Config.get_default_model_path() 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: モデルファイルのパス（Noneの場合はデフォルトパスを使用） """ if hasattr(self, '_initialized') and self._initialized: return self.model_path = model_path or Config.get_default_model_path() 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) -> 'AI': """モデルインスタンスを取得（常駐キャッシュから）""" 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"[AI] モデルパスが存在しません: {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("[AI] GPU検出: CUDAを使用します") else: print("[AI] GPU未検出: CPUモードで実行します") print(f"[AI] モデルをロード中: {model_path}") print(f"[AI] デバイス: {device}") # モデルパスがリポジトリID（"user/repo"形式）か、ローカルパスかを判定 hf_token = os.getenv("HF_TOKEN") is_repo_id = "/" in model_path and not os.path.exists(model_path) if is_repo_id: print(f"[AI] Hugging Face Hub から直接読み込み: {model_path}") else: print(f"[AI] ローカルパスから読み込み: {model_path}") # トークナイザーとモデルをロード（Hubから直接読み込む） tokenizer = AutoTokenizer.from_pretrained( model_path, token=hf_token, ) model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.float16 if device == "cuda" else torch.float32, device_map="auto" if device == "cuda" else None, token=hf_token, ) if device == "cpu": model = model.to(device) # モデルとトークナイザーをタプルで返す print(f"[AI] モデルロード成功 ({device}モード)") return (model, tokenizer) except Exception as e: import traceback print(f"[AI] transformersでのロードに失敗: {e}") traceback.print_exc() return None except Exception as e: import traceback print(f"[AI] モデルロードエラー: {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の特殊トークンを定義 LLAMA_SPECIAL_TOKENS = [ "<|begin_of_text|>", "<|end_of_text|>", "<|eot_id|>", "<|start_header_id|>", "<|end_header_id|>", ] def _clean_text(text: str) -> str: """制御文字・不可視文字・置換文字を厳密に取り除く（正規タグは保持）""" if not text: return "" # 制御文字（0x00-0x1F、0x7F-0x9F）を除去 # ただし、改行・タブ・復帰は許可 cleaned = [] for ch in text: code = ord(ch) # 許可する制御文字: 改行(0x0A), タブ(0x09), 復帰(0x0D) if code in [0x09, 0x0A, 0x0D]: cleaned.append(ch) # 通常の印刷可能文字（0x20-0x7E、およびその他のUnicode印刷可能文字） elif ch.isprintable(): # 置換文字（U+FFFD）を除去 if ch != "\uFFFD": cleaned.append(ch) # その他の制御文字や不可視文字は除去 result = "".join(cleaned) # ゼロ幅文字を除去 result = result.replace("\u200B", "") # Zero-width space result = result.replace("\u200C", "") # Zero-width non-joiner result = result.replace("\u200D", "") # Zero-width joiner result = result.replace("\uFEFF", "") # Zero-width no-break space # その他の不可視文字（結合文字など）を除去 result = result.replace("\u200E", "") # Left-to-right mark result = result.replace("\u200F", "") # Right-to-left mark result = result.replace("\u202A", "") # Left-to-right embedding result = result.replace("\u202B", "") # Right-to-left embedding result = result.replace("\u202C", "") # Pop directional formatting result = result.replace("\u202D", "") # Left-to-right override result = result.replace("\u202E", "") # Right-to-left override return result.strip() for idx, prob in zip(top_indices, top_probs): token_id = idx.item() # 正規タグ（<|eot_id|>など）を保持するため、skip_special_tokens=False token = tokenizer.decode([token_id], skip_special_tokens=False, clean_up_tokenization_spaces=False) token = _clean_text(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("モデルがサポートされていません") return [] except Exception as e: print(f"トークン確率取得エラー: {e}") import traceback traceback.print_exc() return [] def _softmax_from_logprobs(self, logprobs: List[float]) -> List[float]: """logprobsをsoftmaxで確率に変換""" if not logprobs: return [] # 数値安定性のため最大値を引く max_logprob = max(logprobs) exp_logprobs = [exp(logprob - max_logprob) for logprob in logprobs] sum_exp = sum(exp_logprobs) if sum_exp == 0: return [0.0] * len(logprobs) return [exp_logprob / sum_exp for exp_logprob in exp_logprobs] def exp(x: float) -> float: """指数関数の近似実装（math.expの代替）""" import math return math.exp(x)