61 lines
1.2 KiB
Python
61 lines
1.2 KiB
Python
|
import logging
|
||
|
|
from ..config import settings
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
_llm = None
|
||
|
|
|
||
|
|
|
||
|
|
def _load_model():
|
||
|
|
global _llm
|
||
|
|
if _llm is not None:
|
||
|
|
return
|
||
|
|
|
||
|
|
from llama_cpp import Llama
|
||
|
|
|
||
|
|
logger.info("Loading TranslateGemma 12B from %s", settings.translate_model_path)
|
||
|
|
_llm = Llama(
|
||
|
|
model_path=settings.translate_model_path,
|
||
|
|
n_gpu_layers=settings.translate_n_gpu_layers,
|
||
|
|
n_ctx=2048,
|
||
|
|
verbose=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def translate(text: str) -> str:
|
||
|
|
_load_model()
|
||
|
|
|
||
|
|
if not text.strip():
|
||
|
|
return ""
|
||
|
|
|
||
|
|
lines = [
|
||
|
|
"Translate the following English text to Danish.",
|
||
|
|
"Preserve the original meaning, tone, and style.",
|
||
|
|
"Output only the Danish translation, nothing else.",
|
||
|
|
"",
|
||
|
|
f"English: {text}",
|
||
|
|
"",
|
||
|
|
"Danish:",
|
||
|
|
]
|
||
|
|
prompt = "\n".join(lines)
|
||
|
|
|
||
|
|
stop_tokens = ["\n\n", "English:", "User:"]
|
||
|
|
output = _llm(
|
||
|
|
prompt,
|
||
|
|
max_tokens=1024,
|
||
|
|
temperature=0.1,
|
||
|
|
stop=stop_tokens,
|
||
|
|
echo=False,
|
||
|
|
)
|
||
|
|
|
||
|
|
result = output["choices"][0]["text"].strip()
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def unload_model():
|
||
|
|
global _llm
|
||
|
|
_llm = None
|
||
|
|
import torch
|
||
|
|
|
||
|
|
torch.cuda.empty_cache()
|