From 92cdadbfdfb5a5db7cf6f96469ad4370fb3088bd Mon Sep 17 00:00:00 2001
From: Risin <5331402+risin42@users.noreply.github.com>
Date: Tue, 15 Oct 2024 18:49:04 +0900
Subject: [PATCH] modify: model settings

---
 book_maker/translator/gemini_translator.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/book_maker/translator/gemini_translator.py b/book_maker/translator/gemini_translator.py
index 119114e..fc4eb61 100644
--- a/book_maker/translator/gemini_translator.py
+++ b/book_maker/translator/gemini_translator.py
@@ -16,20 +16,15 @@ generation_config = {
     "temperature": 1.0,
     "top_p": 1,
     "top_k": 1,
-    "max_output_tokens": 2048,
+    "max_output_tokens": 8192,
 }
 
-safety_settings = [
-    {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-    {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_MEDIUM_AND_ABOVE"},
-    {
-        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
-    },
-    {
-        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
-        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
-    },
+safety_settings = {
+    "HATE": "BLOCK_NONE",
+    "HARASSMENT": "BLOCK_NONE",
+    "SEXUAL": "BLOCK_NONE",
+    "DANGEROUS": "BLOCK_NONE",
+}
 
 PROMPT_ENV_MAP = {
     "user": "BBM_GEMINIAPI_USER_MSG_TEMPLATE",