File size: 1,133 Bytes
716f8d1 fca9809 716f8d1 7f191bc 716f8d1 ebbd5df 6f13b8c 7f191bc 716f8d1 7f191bc 716f8d1 fca9809 716f8d1 6f13b8c ebbd5df 7f191bc abe0fe2 716f8d1 7f191bc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 |
from logging import getLogger
from llama_cpp import Llama
from functools import lru_cache
logger = getLogger(__name__)
class QwenTranslator:
def __init__(self, model_path, system_prompt_en="", system_prompt_zh="") -> None:
self.llm = Llama(
model_path=model_path,
chat_format="chatml",
verbose=False)
self.sys_prompt_en = system_prompt_en
self.sys_prompt_zh = system_prompt_zh
def to_message(self, prompt, src_lang, dst_lang):
"""构造提示词"""
return [
{"role": "system", "content": self.sys_prompt_en if src_lang == "en" else self.sys_prompt_zh},
{"role": "user", "content": prompt},
]
@lru_cache(maxsize=10)
def translate(self, prompt, src_lang, dst_lang) -> str:
message = self.to_message(prompt, src_lang, dst_lang)
output = self.llm.create_chat_completion(messages=message, temperature=0)
return output['choices'][0]['message']['content']
def __call__(self, prompt,*args, **kwargs):
return self.llm(
prompt,
*args,
**kwargs
)
|