55 lines
1.9 KiB
Python
55 lines
1.9 KiB
Python
|
import asyncio
|
||
|
|
from concurrent.futures import ThreadPoolExecutor
|
||
|
|
from dotenv import find_dotenv, load_dotenv
|
||
|
|
from ollama import generate
|
||
|
|
from Craft.module.ml.prompt import generate_system, generate_user, generate_system_eng, generate_user_eng
|
||
|
|
from Craft.module.ml.util import extract_emoji_and_text
|
||
|
|
|
||
|
|
load_dotenv(find_dotenv())
|
||
|
|
|
||
|
|
class Engine:
|
||
|
|
def __init__(self):
|
||
|
|
self.key = None
|
||
|
|
|
||
|
|
def _generate(self, first_word: str, second_word: str, eng_result: str) -> dict:
|
||
|
|
return generate(
|
||
|
|
model="infcraft:latest",
|
||
|
|
system=generate_system(first_word, second_word, eng_result),
|
||
|
|
prompt=generate_user(first_word, second_word),
|
||
|
|
keep_alive=60*60*24,
|
||
|
|
context=None,
|
||
|
|
options={
|
||
|
|
"seed": 0,
|
||
|
|
"temperature": 0.4,
|
||
|
|
"top_p": 0.85,
|
||
|
|
"top_k": 0.1,
|
||
|
|
"max_tokens": 32,
|
||
|
|
"main_gpu": 0,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
async def _generate_eng(self, first_word: str, second_word: str) -> dict:
|
||
|
|
system = await generate_system_eng(first_word, second_word)
|
||
|
|
prompt = await generate_user_eng(first_word, second_word)
|
||
|
|
return generate(
|
||
|
|
model="mistral:latest",
|
||
|
|
system=system,
|
||
|
|
prompt=prompt,
|
||
|
|
keep_alive=60*60*24,
|
||
|
|
options={
|
||
|
|
"seed": 0,
|
||
|
|
"temperature": 0.8,
|
||
|
|
"top_p": 1,
|
||
|
|
"top_k": 0.1,
|
||
|
|
"max_tokens": 64,
|
||
|
|
"main_gpu": 0,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
async def generate(self, first_word: str, second_word: str) -> str:
|
||
|
|
loop = asyncio.get_running_loop()
|
||
|
|
eng_result = await self._generate_eng(first_word, second_word)
|
||
|
|
with ThreadPoolExecutor() as executor:
|
||
|
|
data = await loop.run_in_executor(executor, self._generate, first_word, second_word, eng_result['response'])
|
||
|
|
return extract_emoji_and_text(data['response'])
|