Merge pull request #2 from Masao-Someki/feature/add_gemini

modules/llm/__init__.py

@@ -2,10 +2,11 @@ from .base import AbstractLLMModel
 from .registry import LLM_MODEL_REGISTRY, get_llm_model, register_llm_model
 from .hf_pipeline import HFTextGenerationLLM
 from .qwen import QwenLLM
+from .gemini import GeminiLLM
 
 __all__ = [
     "AbstractLLMModel",
     "get_llm_model",
     "register_llm_model",
     "LLM_MODEL_REGISTRY",
-]
+]

modules/llm/gemini.py

@@ -0,0 +1,31 @@
+import os
+
+os.environ["XDG_CACHE_HOME"] = "./.cache" # must be set before importing google.generativeai
+
+import google.generativeai as genai
+
+from .base import AbstractLLMModel
+from .registry import register_llm_model
+
+
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+
+@register_llm_model("gemini-2.5-flash")
+class GeminiLLM(AbstractLLMModel):
+    def __init__(
+        self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
+    ):
+        if os.environ.get("XDG_CACHE_HOME") != cache_dir:
+            raise RuntimeError(
+                f"XDG_CACHE_HOME must be set to '{cache_dir}' before importing this module."
+            )
+        if not GOOGLE_API_KEY:
+            raise ValueError("Please set the GOOGLE_API_KEY environment variable to use Gemini.")
+        super().__init__(model_id=model_id, **kwargs)
+        genai.configure(api_key=GOOGLE_API_KEY)
+        self.model = genai.GenerativeModel(model_id)
+
+    def generate(self, prompt: str, **kwargs) -> str:
+        response = self.model.generate_content(prompt, **kwargs)
+        return response.text

requirements.txt

@@ -12,9 +12,9 @@ pykakasi
 basic-pitch[onnx]
 audiobox_aesthetics
 transformers
-s3prl
 zhconv
 git+https://github.com/sea-turt1e/kanjiconv
 soundfile
 PyYAML
 gradio
+google-generativeai

tests/test_llm_infer.py

@@ -11,6 +11,7 @@ if __name__ == "__main__":
         # "deepseek-ai/DeepSeek-R1-0528",
         # "openai-community/gpt2-xl",
         # "google/gemma-2-2b",
+        # "gemini-2.5-flash",
     ]
     for model_id in supported_llms:
         try: