import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, StoppingCriteriaList
import os
import torch

theme = "darkgrass"
title = "Polyglot(Korean) Demo"
model_name = "EleutherAI/polyglot-ko-1.3b"

bad_words = [
  '...',
  '....',
  '(중략)',
  'http'
]

description = "polyglot (1.3B 파라미터 사이즈) 한국어 모델을 시연하는 데모페이지 입니다."
article = "<p style='text-align: center'><a href='https://github.com/EleutherAI/polyglot' target='_blank'>Polyglot: Large Language Models of Well-balanced Competence in Multi-languages</a></p>"
examples = [
    ["CPU와 GPU의 차이는,"],
    ["질문: 우크라이나 전쟁이 세계3차대전으로 확전이 될까요? \n답변:"],
    ["2040년 미국은, "]
]
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
  model_name
)
model.eval()

pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, eos_token_id=tokenizer.eos_token_id)

def predict(text):
  with torch.no_grad():
    tokens = tokenizer(text, return_tensors="pt").input_ids
    # generate and end generate if <|endoftext|> is not in text
    gen_tokens = model.generate(
      tokens, do_sample=True, temperature=0.8, max_new_tokens=64, top_k=50, top_p=0.8,
      no_repeat_ngram_size=3, repetition_penalty=1.2,
      bad_words_ids=[
        tokenizer.encode(bad_word) for bad_word in bad_words
      ],
      eos_token_id=tokenizer.eos_token_id,
      pad_token_id=tokenizer.pad_token_id
    )
    generated = tokenizer.batch_decode(gen_tokens)[0]
    return generated
  # return pipe(text)[0]['generated_text']

iface = gr.Interface(
  fn=predict, 
  inputs='text',
  outputs='text',
  examples=examples
)

iface.launch()





  
# print(generated)  # print: 인간처럼 생각하고, 행동하는 '지능'을 통해 인류가 이제까지 풀지 못했던 문제의 해답을 찾을 수 있을 것이다. 과학기술이 고도로 발달한 21세기를 살아갈 우리 아이들에게 가장 필요한 것은 사고력 훈련이다. 사고력 훈련을 통해, 세상