要約機能の実装

src/ai_api/config.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    """使用するAIモデルに関する情報を一元管理します。"""
+
+    NAME: str = "llm-jp/t5-small-japanese-finetuned-sum"
+    REVISION: str = "main"

src/ai_api/core/inference.py

@@ -1,6 +1,28 @@
+from typing import cast
+
+from transformers import pipeline
+
+from ai_api.config import ModelConfig
+
+
 class Summarizer:
     """
     AIモデルを管理し、テキスト要約を実行するクラス。
     """
 
-    pass
+    def __init__(self, config: ModelConfig) -> None:
+        self.config = config
+        # __init__で一度だけpipelineを初期化
+        self.summarizer = pipeline(
+            "summarization",
+            model=self.config.NAME,
+            revision=self.config.REVISION,
+        )
+
+    def summarize(self, text: str) -> str:
+        """
+        与えられたテキストを要約する。
+        """
+        # 保持しているsummarizerを使って要約
+        result = self.summarizer(text)
+        return cast(str, result[0]["summary_text"])

tests/core/test_inference.py

@@ -1,8 +1,33 @@
+from ai_api.config import ModelConfig
+from ai_api.core.inference import Summarizer
 
 
-def test_summarizer_can_be_imported() -> None:
+def test_summarizer_initialization_with_test_model() -> None:
     """
-    Summarizerクラスがインポートできることをテストする。
+    テスト用の軽量モデルでSummarizerが初期化できることをテストする。
     """
-    # このテストは、Summarizerクラスがインポートできれば成功する
-    assert True
+    # Arrange: テスト専用の軽量モデルを指定
+    config = ModelConfig(NAME="sshleifer/distilbart-cnn-6-6", REVISION="main")
+
+    # Act: 実際にモデルをロードして初期化
+    summarizer = Summarizer(config=config)
+
+    # Assert
+    assert isinstance(summarizer, Summarizer)
+
+
+def test_summarize_with_test_model() -> None:
+    """
+    テスト用の軽量モデルでsummarizeメソッドが動作することをテストする。
+    """
+    # Arrange: テスト専用の軽量モデルを指定
+    config = ModelConfig(NAME="sshleifer/distilbart-cnn-6-6", REVISION="main")
+    summarizer = Summarizer(config=config)
+    text = "This is a test sentence. It is a very nice sentence to summarize."
+
+    # Act: 実際に要約を実行
+    summary = summarizer.summarize(text)
+
+    # Assert: 要約結果が文字列であり、空でないことを確認
+    assert isinstance(summary, str)
+    assert len(summary) > 0