mirror of
https://github.com/yihong0618/bilingual_book_maker.git
synced 2025-06-06 19:45:30 +00:00
6c72c29943
* Minor improvements for loader and translator * use rotate key * fix: typo --------- Co-authored-by: yihong0618 <zouzou0208@gmail.com>
42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
import requests
|
||
from rich import print
|
||
|
||
from .base_translator import Base
|
||
|
||
|
||
class GPT3(Base):
|
||
def __init__(self, key, language, api_base=None):
|
||
super().__init__(key, language)
|
||
self.api_url = (
|
||
f"{api_base}v1/completions"
|
||
if api_base
|
||
else "https://api.openai.com/v1/completions"
|
||
)
|
||
self.headers = {
|
||
"Content-Type": "application/json",
|
||
}
|
||
# TODO support more models here
|
||
self.data = {
|
||
"prompt": "",
|
||
"model": "text-davinci-003",
|
||
"max_tokens": 1024,
|
||
"temperature": 1,
|
||
"top_p": 1,
|
||
}
|
||
self.session = requests.session()
|
||
self.language = language
|
||
|
||
def rotate_key(self):
|
||
self.headers["Authorization"] = f"Bearer {next(self.keys)}"
|
||
|
||
def translate(self, text):
|
||
print(text)
|
||
self.rotate_key()
|
||
self.data["prompt"] = f"Please help me to translate，`{text}` to {self.language}"
|
||
r = self.session.post(self.api_url, headers=self.headers, json=self.data)
|
||
if not r.ok:
|
||
return text
|
||
t_text = r.json().get("choices")[0].get("text", "").strip()
|
||
print(t_text)
|
||
return t_text
|