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

theme = "darkgrass"
title = "GPT-NeoX(Korean) Demo"
model_name = "EleutherAI/gpt-neox-ko-1.3b"

description = "GPT-NeoX 한국어 모델을 시연하는 데모페이지 입니다."
# article = "<p style='text-align: center'><a href='https://github.com/kingoflolz/mesh-transformer-jax' target='_blank'>GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model</a></p>"
examples = [
    ["인간처럼 생각하고, 행동하는 '지능'을 통해"],
    ["질문: 한가위를 가장 의미있게 보낼수 있는 방법은 무엇일까요? \n답변:"],
    ["2040년 미국은, "]
]

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-ko-1.3b", use_auth_token=os.environ['TOKEN'])
model = AutoModelForCausalLM.from_pretrained(
  "EleutherAI/gpt-neox-ko-1.3b", use_auth_token=os.environ['TOKEN'],
  pad_token_id=tokenizer.pad_token_id,
  eos_token_id=tokenizer.eos_token_id,
  bos_token_id=tokenizer.bos_token_id  
)
model.eval()

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

def predict(text):
  stopping_cond = StoppingCriteriaList([tokenizer.encode('<|endoftext|>')])
  with torch.no_grad():
    tokens = tokenizer(text, return_tensors="pt").input_ids
    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  
    )
    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세기를 살아갈 우리 아이들에게 가장 필요한 것은 사고력 훈련이다. 사고력 훈련을 통해, 세상