Rewrote gemini.py to utilize context and system prompt (#801)

* Fix Google Vision path to api key

* Rewrote gemini.py to utilize context and system prompt

Rewrote gemini.py to utilize context and system prompt. Needs further testing as you can't switch to other translators (like chatgpt) from gemini
unless you restart the program.

* Update translatorsetting.json

* Update gemini.py
This commit is contained in:
blastentwice 2024-06-06 04:47:18 -07:00 committed by GitHub
parent 4bc6a1ae0d
commit fc1493eb77
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 105 additions and 43 deletions

View File

@ -1,5 +1,5 @@
from translator.basetranslator import basetrans
import re
class TS(basetrans):
@ -21,47 +21,82 @@ class TS(basetrans):
"ar": "Arabic",
"th": "Thai",
}
def __init__(self, typename):
self.context = []
super().__init__(typename)
def translate(self, content):
if self.config["使用自定义promt"]:
prompt = self.config["自定义promt"]
else:
prompt = "You are a translator. Please help me translate the following {} text into {}, and you should only tell me the translation.".format(
self.srclang, self.tgtlang
)
res = self.session.post(
"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent",
params={"key": self.config["api-key"]},
json={
"contents": [
{"role": "user", "parts": [{"text": prompt + "\n" + content}]}
],
"generation_config": {"candidate_count": 1, "temperature": 0.4},
"safety_settings": [
{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE"},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE",
},
],
def inittranslator(self):
self.api_key = None
def translate(self, query):
self.checkempty(["SECRET_KEY", "model"])
self.contextnum = int(self.config["context"])
try:
gen_config = {"generationConfig": {"stopSequences": [" \n"], "temperature": float(self.config["Temperature"])}}
except:
gen_config = {"generationConfig": {"temperature": float(0.3)}}
try:
model= self.config["model"]
except:
model= "gemini-1.5-flash"
safety={"safety_settings": [
{"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE",
},
verify=False,
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE",
}
]
}
if self.config["use_custom_prompt"]:
sys_message = {"systemInstruction": {"parts": {"text":self.config["custom_prompt"]}}}
else:
sys_message= {"systemInstruction":
{
"parts": {
"text":"You are a translator. Please help me translate the following {} text into {}, and you should only tell me the translation.".format(
self.srclang, self.tgtlang
),
},
},
}
message=[]
for _i in range(min(len(self.context) // 2, self.contextnum)):
i = (
len(self.context) // 2
- min(len(self.context) // 2, self.contextnum)
+ _i
)
message.append(self.context[i * 2])
message.append(self.context[i * 2 + 1])
message.append({"role": "user", "parts":[{"text":query}]})
contents=dict(contents=message)
payload= {**contents, **safety, **sys_message, **gen_config }
res = self.session.post(
f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent",
params={"key": self.config["SECRET_KEY"]},
json=payload
)
try:
res = res.json()
except:
raise Exception(res.text)
try:
line = res["candidates"][0]["content"]["parts"][0]["text"]
return line
line = res.json()["candidates"][0]["content"]["parts"][0]["text"]
yield line
except:
print(res)
raise Exception("Error")
self.context.append({"role": "user", "parts": [{"text": query}]})
self.context.append({"role": "model", "parts": [{"text": line}]})

View File

@ -470,17 +470,44 @@
"gemini": {
"args": {
"注册网址": "https://ai.google.dev/tutorials/python_quickstart",
"api-key": "",
"使用自定义promt": false,
"自定义promt": ""
"注册网址2": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models",
"SECRET_KEY": "",
"Temperature": 0.3,
"model": "gemini-1.5-flash",
"context": 0,
"use_custom_prompt": false,
"custom_prompt": "",
},
"argstype": {
"custom_prompt": {
"name": "自定义promt"
},
"use_custom_prompt": {
"type": "switch",
"name": "使用自定义promt"
},
"context": {
"type": "intspin",
"min": 0,
"max": 10,
"step": 1,
"name": "附带上下文个数"
},
"注册网址": {
"type": "label",
"islink": true
},
"使用自定义promt": {
"type": "switch"
"注册网址2": {
"type": "label",
"islink": true
},
"Temperature": {
"type": "spin",
"min": 0,
"max": 1,
"step": 0.1,
}
}
},