54 lines
1.9 KiB
Python
54 lines
1.9 KiB
Python
import asyncio
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dotenv import find_dotenv, load_dotenv
|
|
from ollama import generate, AsyncClient
|
|
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
|
|
self.client = AsyncClient(host="http://100.64.0.49:7020")
|
|
|
|
async def _generate(self, first_word: str, second_word: str) -> dict:
|
|
return await self.client.generate(
|
|
model="mergegame:latest",
|
|
system=generate_system(first_word, second_word),
|
|
prompt=generate_user(first_word, second_word),
|
|
keep_alive=60*60*24,
|
|
context=None,
|
|
options={
|
|
"seed": 0,
|
|
"temperature": 0.01,
|
|
"top_p": 1,
|
|
"top_k": 0,
|
|
"max_tokens": 16,
|
|
"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:
|
|
#eng_result = await self._generate_eng(first_word, second_word)
|
|
data = await self._generate(first_word, second_word)
|
|
return extract_emoji_and_text(data['response'])
|