ms180 commited on
Commit
4b87b66
·
1 Parent(s): be053b4

Add Gemini

Browse files
Files changed (2) hide show
  1. modules/llm/gemini.py +23 -0
  2. requirements.txt +1 -1
modules/llm/gemini.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import google.generativeai as genai
2
+ from abc import ABC, abstractmethod
3
+ from .base import AbstractLLMModel
4
+ from .registry import register_llm_model
5
+
6
+
7
+
8
+ GEMINI_TOKEN = os.getenv("GEMINI_API_KEY")
9
+
10
+
11
+ @register_llm_model("gemini-2.5-flash")
12
+ class GeminiModel(AbstractLLMModel):
13
+ def __init__(
14
+ self, model_id: str, device: str = "cpu", cache_dir: str = "cache", **kwargs
15
+ ):
16
+ super().__init__(model_id=model_id, **kwargs)
17
+ genai.configure(api_key=GEMINI_API_KEY)
18
+ self.model = genai.GenerativeModel(model_id)
19
+
20
+ def generate(self, prompt: str, **kwargs) -> str:
21
+ response = self.model.generate_content(prompt, **kwargs)
22
+ return response.text
23
+
requirements.txt CHANGED
@@ -12,9 +12,9 @@ pykakasi
12
  basic-pitch[onnx]
13
  audiobox_aesthetics
14
  transformers
15
- s3prl
16
  zhconv
17
  git+https://github.com/sea-turt1e/kanjiconv
18
  soundfile
19
  PyYAML
20
  gradio
 
 
12
  basic-pitch[onnx]
13
  audiobox_aesthetics
14
  transformers
 
15
  zhconv
16
  git+https://github.com/sea-turt1e/kanjiconv
17
  soundfile
18
  PyYAML
19
  gradio
20
+ google-generativeai