kocraft-backend/test.py

67 lines
2.2 KiB
Python
Raw Normal View History

2024-07-15 08:47:42 +00:00
import os
import re
import asyncio
import emoji
import aiogoogletrans
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())
translator = aiogoogletrans.Translator()
class Engine:
def __init__(self):
self.key = None
def _generate(self, first_word: str, second_word: str, eng_result: str) -> str:
gen_data = generate(
model="infcraft:minial",
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.2,
"top_p" : 0.85,
"top_k" : 0.1,
"max_tokens" : 32,
"main_gpu" : 0,
}
)
# 반환된 값의 response에서 이모지만 추출
return gen_data
async def _generate_eng(self, first_word: str, second_word: str) -> str:
gen_data = generate(
model="mistral:latest",
system=await generate_system_eng(first_word, second_word),
prompt=await generate_user_eng(first_word, second_word),
keep_alive=60*60*24,
options={
"seed" : 0,
"temperature" : 0.2,
"top_p" : 1,
"top_k" : 0.1,
"max_tokens" : 64,
"main_gpu" : 0,
}
)
return gen_data
async def generate(self, first_word: str, second_word: str) -> str:
with ThreadPoolExecutor() as executor:
loop = asyncio.get_event_loop()
eng_result = await self._generate_eng(first_word, second_word)
data = await loop.run_in_executor(executor, self._generate, first_word, second_word, eng_result['response'])
return data
data = asyncio.run(Engine().generate("👮 윤석열", "👮 이재명"))
print(extract_emoji_and_text(data['response']))