from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model_id = "heegyu/42dot_LLM-PLM-1.3B-mt"
model = AutoModelForCausalLM.from_pretrained(model_id).eval().half()
tokenizer = AutoTokenizer.from_pretrained(model_id)

if torch.cuda.is_available():
  device = "cuda:0"
  model.to(device)
else:
  device = "cpu"

@torch.no_grad()
def generate_text(prompt):
  input_ids = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False).to(device)
  output_ids = model.generate(input_ids, min_new_tokens=4, max_length=1024, early_stopping=True)
  output_ids = output_ids.cpu()[0][len(input_ids[0]):]
  print(tokenizer.decode(output_ids))

bos, eos = tokenizer.bos_token, tokenizer.eos_token

# 한 -> 영
text = "삼성전자가 갤럭시 스마트폰·태블릿 전용 ‘클라우드 게임 플랫폼’을 이르면 이달 공개한다. 전 세계 10억 명의 갤럭시 사용자가 콘솔(게임기)을 구매하거나 게임 앱을 내려받지 않아도 스마트폰을 통해 실시간으로 유명 게임을 즐길 수 있게 되는 것이다. 기기 판매에 의존하지 않고 안정적인 서비스 수익을 올리려는 삼성전자의 ‘신사업 승부수’란 평가가 나온다."
generate_text(f"{bos} {text} {eos} ")
# ㈜Samsung Electronics will release the Cloud Game Platform for Galaxy smartphones and tablets in early this month, allowing users of 1 billion people around the world to enjoy famous games on their smartphones in real time without buying consoles or downloading game apps. It is said to be a 'business move' by Samsung Electronics, which is trying to earn stable service revenue without relying on sales.<|endoftext|>

# 영 -> 한, 마지막 문장 짤렸음
text = """Samsung Electronics will unveil a "cloud game platform" exclusively for Galaxy smartphones and tablets as early as this month. One billion Galaxy users around the world will be able to enjoy famous games in real time through smartphones without having to purchase consoles or download game apps. Analysts say that Samsung Electronics is a "new business winning move" to earn stable service profits without relying on device sales."""
generate_text(f"{bos} {text} {eos} ")
# ＮＣ는 이달 중 갤럭시 스마트폰과 태블릿 전용 '클라우드 게임 플랫폼'을 독점 공개할 예정인데, 전 세계 1억명의 갤럭시 사용자들은 콘솔이나 게임 앱 다운로드 없이 스마트폰을 통해 유명 게임을 실시간으로 즐길 수 있게 됐다.<|endoftext|>

# 영 -> 한, 앞에 번역한 단어를 지정할 수 있다.
text = """Samsung Electronics will unveil a "cloud game platform" exclusively for Galaxy smartphones and tablets as early as this month."""
generate_text(f"{bos} Samsung Electronics {eos} 삼성전자 {eos} {bos} {text} {eos} ")
# Ｎ가전 삼성전자가 갤럭시 스마트폰과 태블릿 전용 '클라우드 게임 플랫폼'을 이달 중으로 공개한다.<|endoftext|>

python main.py \
    --model hf-causal \
    --model_args pretrained=heegyu/42dot_LLM-PLM-1.3B-mt \
    --tasks kobest_hellaswag,kobest_copa,kobest_boolq,kobest_sentineg \
    --device cuda:0