Switch from llama-cpp-python to transformers for PyTorch model support

app.py

@@ -40,10 +40,7 @@ adapter = None
 status_message = "モデル初期化中..."
 status_lock = threading.Lock()
 
-HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "WatNeru/LLMView-model")
-HF_MODEL_FILENAME = os.getenv(
-    "HF_MODEL_FILENAME", "llama-3.2-3b-instruct-q4_k_m.gguf"
-)
+HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "meta-llama/Llama-3.2-3B-Instruct")
 HF_LOCAL_DIR = Path(
     os.getenv(
         "HF_MODEL_LOCAL_DIR",
@@ -72,25 +69,27 @@ def _set_status(message: str) -> None:
 
 
 def ensure_model_available() -> str:
-    """モデルファイルをローカルに用意（なければHFから取得）"""
+    """モデルディレクトリをローカルに用意（なければHFから取得）"""
     print(f"[MODEL] ensure_model_available() 開始")
     current_path = Path(path_manager.get_model_path())
     print(f"[MODEL] 現在のモデルパス: {current_path}")
     
-    if current_path.exists():
-        print(f"[MODEL] 既存のモデルファイルを使用: {current_path}")
-        return str(current_path)
-
-    print(f"[MODEL] モデルファイルが見つからないため、ダウンロードを開始")
+    # PyTorch モデルの場合、ディレクトリ全体をチェック
+    if current_path.exists() and current_path.is_dir():
+        # config.json があるか確認（モデルディレクトリの確認）
+        if (current_path / "config.json").exists():
+            print(f"[MODEL] 既存のモデルディレクトリを使用: {current_path}")
+            return str(current_path)
+
+    print(f"[MODEL] モデルディレクトリが見つからないため、ダウンロードを開始")
     HF_LOCAL_DIR.mkdir(parents=True, exist_ok=True)
     print(f"[MODEL] ダウンロード先ディレクトリ: {HF_LOCAL_DIR}")
     _set_status("Hugging Face からモデルをダウンロード中...")
 
-    # snapshot_downloadの戻り値（ダウンロード先ディレクトリ）を使用
+    # snapshot_downloadでモデル全体をダウンロード（PyTorch モデルは複数ファイル）
     try:
         downloaded_dir = snapshot_download(
             repo_id=HF_MODEL_REPO,
-            allow_patterns=HF_MODEL_FILENAME,
             local_dir=str(HF_LOCAL_DIR),
             local_dir_use_symlinks=False,
             token=HF_TOKEN,
@@ -102,32 +101,17 @@ def ensure_model_available() -> str:
         traceback.print_exc()
         raise
     
-    # ダウンロードされたディレクトリからファイルを検索
+    # ダウンロードされたディレクトリを確認
     downloaded_dir_path = Path(downloaded_dir)
     print(f"[MODEL] ダウンロード先パス: {downloaded_dir_path}")
     
-    # まず直接パスを試す
-    downloaded = downloaded_dir_path / HF_MODEL_FILENAME
-    print(f"[MODEL] 直接パスを確認: {downloaded}")
-    
-    if not downloaded.exists():
-        # リポジトリ構造を保持している可能性があるので、再帰的に検索
-        print(f"[MODEL] 直接パスに存在しないため、再帰的に検索中...")
-        found_files = list(downloaded_dir_path.rglob(HF_MODEL_FILENAME))
-        print(f"[MODEL] 見つかったファイル数: {len(found_files)}")
-        if found_files:
-            downloaded = found_files[0]
-            print(f"[MODEL] ファイルを発見: {downloaded}")
-        else:
-            # ディレクトリ内の全ファイルをリストアップしてデバッグ情報を出力
-            all_files = list(downloaded_dir_path.rglob("*"))
-            print(f"[MODEL] ディレクトリ内の全ファイル: {[str(f) for f in all_files[:20]]}")
-            raise FileNotFoundError(
-                f"モデル {HF_MODEL_FILENAME} が {downloaded_dir} に見つかりません。"
-                f"見つかったファイル: {[str(f) for f in all_files[:10]]}"
-            )
-
-    model_path_str = str(downloaded.resolve())
+    # config.json があるか確認
+    if not (downloaded_dir_path / "config.json").exists():
+        raise FileNotFoundError(
+            f"モデルディレクトリ {downloaded_dir} に config.json が見つかりません。"
+        )
+
+    model_path_str = str(downloaded_dir_path.resolve())
     print(f"[MODEL] モデルパスを設定: {model_path_str}")
     os.environ["LLM_MODEL_PATH"] = model_path_str
     path_manager.model_path = model_path_str

package/ai.py

@@ -47,64 +47,47 @@ class AI:
         cls._instances.clear()
     
     def _load_model(self, model_path: str) -> Optional[Any]:
-        """モデルをロード"""
+        """モデルをロード（Transformers使用）"""
         try:
             if not model_path or not os.path.exists(model_path):
                 return None
                 
-            # llama-cpp-pythonを使用してモデルをロード
+            # transformersを使用してモデルをロード
             try:
-                from llama_cpp import Llama
+                from transformers import AutoModelForCausalLM, AutoTokenizer
+                import torch
                 
                 # GPUが利用可能かチェック
-                use_gpu = False
-                try:
-                    # CUDAが利用可能かチェック
-                    result = subprocess.run(
-                        ["nvidia-smi"], 
-                        capture_output=True, 
-                        text=True, 
-                        timeout=2
-                    )
-                    if result.returncode == 0:
-                        # GPUが利用可能
-                        use_gpu = True
-                        print("[AI] GPU検出: CUDAを使用します")
-                    else:
-                        print("[AI] GPU未検出: CPUモードで実行します")
-                except (subprocess.TimeoutExpired, FileNotFoundError, Exception):
-                    # nvidia-smiが使えない場合はCPUモード
+                device = "cuda" if torch.cuda.is_available() else "cpu"
+                if device == "cuda":
+                    print("[AI] GPU検出: CUDAを使用します")
+                else:
                     print("[AI] GPU未検出: CPUモードで実行します")
                 
-                # GPUモードで試行
-                if use_gpu:
-                    try:
-                        llm = Llama(
-                            model_path=model_path,
-                            n_ctx=2048,
-                            logits_all=True,
-                            n_gpu_layers=-1,  # すべてのレイヤーをGPUに配置
-                            verbose=True,
-                        )
-                        print("[AI] モデルロード成功 (GPUモード)")
-                        return llm
-                    except Exception as gpu_error:
-                        print(f"[AI] GPUモードでのロードに失敗、CPUモードにフォールバック: {gpu_error}")
-                        use_gpu = False  # CPUモードにフォールバック
+                print(f"[AI] モデルをロード中: {model_path}")
+                print(f"[AI] デバイス: {device}")
                 
-                # CPUモードでロード
-                llm = Llama(
-                    model_path=model_path,
-                    n_ctx=2048,
-                    logits_all=True,
-                    n_gpu_layers=0,
-                    verbose=True,
+                # トークナイザーとモデルをロード
+                tokenizer = AutoTokenizer.from_pretrained(
+                    model_path,
+                    token=os.getenv("HF_TOKEN"),
                 )
-                print("[AI] モデルロード成功 (CPUモード)")
-                return llm
+                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=os.getenv("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] llama-cpp-pythonでのロードに失敗: {e}")
+                print(f"[AI] transformersでのロードに失敗: {e}")
                 traceback.print_exc()
                 return None
                 
@@ -129,35 +112,34 @@ class AI:
             return []
             
         try:
-            # llama-cpp-pythonのcreate_completionを使用
-            if hasattr(self.model, "create_completion"):
-                resp = self.model.create_completion(
-                    prompt=text,
-                    max_tokens=1,
-                    logprobs=k,
-                    temperature=0.0,
-                    echo=False,
-                )
+            # transformers モデルの場合
+            if isinstance(self.model, tuple) and len(self.model) == 2:
+                model, tokenizer = self.model
+                import torch
                 
-                # レスポンスからトークンと確率を抽出
-                items: List[Tuple[str, float]] = []
-                choice = resp.get("choices", [{}])[0]
-                lp = choice.get("logprobs", {})
-                top = lp.get("top_logprobs", [])
+                # テキストをトークン化
+                inputs = tokenizer(text, return_tensors="pt")
+                device = next(model.parameters()).device
+                inputs = {k: v.to(device) for k, v in inputs.items()}
                 
-                if top and isinstance(top[0], dict):
-                    cand_dict = top[0]
-                    tokens = list(cand_dict.keys())
-                    logprobs = [cand_dict[t] for t in tokens]
-                    
-                    # logprobsを確率に変換
-                    probs = self._softmax_from_logprobs(logprobs)
-                    
-                    for token, prob in zip(tokens, probs):
-                        items.append((token, float(prob)))
+                # モデルで推論（勾配計算なし）
+                with torch.no_grad():
+                    outputs = model(**inputs)
+                    logits = outputs.logits[0, -1, :]  # 最後のトークンのlogits
                 
-                # 確率順でソートして上位k個を返す
-                items = sorted(items, key=lambda x: x[1], reverse=True)[:k]
+                # logitsを確率に変換（softmax）
+                probs = torch.softmax(logits, dim=-1)
+                
+                # 上位k個のトークンを取得
+                top_probs, top_indices = torch.topk(probs, k)
+                
+                # トークンIDを文字列に変換
+                items: List[Tuple[str, float]] = []
+                for idx, prob in zip(top_indices, top_probs):
+                    token_id = idx.item()
+                    token = tokenizer.decode([token_id])
+                    prob_value = prob.item()
+                    items.append((token, float(prob_value)))
                 
                 # 確率を正規化
                 if items:
@@ -171,11 +153,13 @@ class AI:
                 
                 return items
             else:
-                print("モデルがcreate_completionメソッドをサポートしていません")
+                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]:

requirements.txt

@@ -10,7 +10,9 @@ sudachipy>=0.6.7
 sudachidict-core>=20240125
 
 # AI/LLM
-llama-cpp-python==0.2.79
+transformers>=4.40.0
+torch>=2.0.0
+accelerate>=0.30.0
 
 # UI
 gradio>=4.38.0