From 696ea341b579992af78498a8370a6a6b7c1ee132 Mon Sep 17 00:00:00 2001
From: Risin <5331402+risin42@users.noreply.github.com>
Date: Tue, 15 Oct 2024 18:37:31 +0900
Subject: [PATCH] feat: add support for --temperature option to gemini

---
 book_maker/cli.py                          |  2 +-
 book_maker/translator/gemini_translator.py | 12 +++++++++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/book_maker/cli.py b/book_maker/cli.py
index f42db4f..d926e0e 100644
--- a/book_maker/cli.py
+++ b/book_maker/cli.py
@@ -290,7 +290,7 @@ So you are close to reaching the limit. You have to choose your own value, there
         "--temperature",
         type=float,
         default=1.0,
-        help="temperature parameter for `chatgptapi`/`gpt4`/`claude`",
+        help="temperature parameter for `chatgptapi`/`gpt4`/`claude`/`gemini`",
     )
     parser.add_argument(
         "--block_size",
diff --git a/book_maker/translator/gemini_translator.py b/book_maker/translator/gemini_translator.py
index 0dc43a6..da8e7f8 100644
--- a/book_maker/translator/gemini_translator.py
+++ b/book_maker/translator/gemini_translator.py
@@ -11,7 +11,7 @@ from rich import print
 from .base_translator import Base
 
 generation_config = {
-    "temperature": 0.7,
+    "temperature": 1.0,
     "top_p": 1,
     "top_k": 1,
     "max_output_tokens": 2048,
@@ -38,9 +38,15 @@ class Gemini(Base):
 
     DEFAULT_PROMPT = "Please help me to translate,`{text}` to {language}, please return only translated content not include the origin text"
 
-    def __init__(self, key, language, **kwargs) -> None:
-        genai.configure(api_key=key)
+    def __init__(
+        self,
+        key,
+        language,
+        temperature=1.0,
+        **kwargs,
+    ) -> None:
         super().__init__(key, language)
+        generation_config["temperature"] = temperature
         model = genai.GenerativeModel(
             model_name="gemini-pro",
             generation_config=generation_config,