import os
os.system("pip install git+https://github.com/huggingface/transformers")

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import torch

tok = AutoTokenizer.from_pretrained("distilgpt2")
model = AutoModelForCausalLM.from_pretrained("distilgpt2")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
n_gpu = 0 if torch.cuda.is_available()==False else torch.cuda.device_count()
model.to(device)

early_stop_pattern = tok.eos_token
print(f'Early stop pattern = \"{early_stop_pattern}\"')

def generate(text = ""):  
  streamer = TextIteratorStreamer(tok, timeout=10.)
  if len(text) == 0:
    text = " "
  inputs = tok([text], return_tensors="pt").to(device)
  generation_kwargs = dict(inputs, streamer=streamer, repetition_penalty=2.0, do_sample=True, top_k=40, top_p=0.97, max_new_tokens=128)
  thread = Thread(target=model.generate, kwargs=generation_kwargs)
  thread.start()
  generated_text = ""
  for new_text in streamer:
    yield generated_text + new_text
    #print(new_text, end ="")
    generated_text += new_text
    if early_stop_pattern in generated_text:
      generated_text = generated_text[: generated_text.find(early_stop_pattern) if early_stop_pattern else None]
      streamer.end()
      #print("\n--\n")
      yield generated_text
      return

demo = gr.Interface(
    title="TextIteratorStreamer + Gradio demo",
    fn=generate,
    inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
    outputs=gr.outputs.Textbox(label="Generated Text"),    
)

demo.queue()
demo.launch()