feat: add safety settings for Gemini translator (#94)

rss-translator · Jul 24, 2024 · 3e4f7d8 · 3e4f7d8
1 parent 07a4f73
commit 3e4f7d8
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/translator/models/gemini.py b/translator/models/gemini.py
@@ -1,5 +1,6 @@
 from config import settings
 import google.generativeai as genai
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 from .base import TranslatorEngine
 import logging
 from time import sleep
@@ -82,7 +83,13 @@ def translate(
                 top_k=self.top_k,
                 max_output_tokens=self.max_tokens,
             )
-            res = model.generate_content(prompt, generation_config=generation_config)
+            safety_settings = {
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+                HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+            }
+            res = model.generate_content(prompt, generation_config=generation_config, safety_settings=safety_settings)
             finish_reason = res.candidates[0].finish_reason if res.candidates else None
             if finish_reason == 1:
                 translated_text = res.text