import torch
import solara
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('gpt2')
model = AutoModelForCausalLM.from_pretrained('gpt2')

prompt = solara.reactive("Alan Turing theorized that computers would one day become")
sample = solara.reactive(False)
answer = solara.reactive("")
seed = solara.reactive(False)
seed_int = solara.reactive(42)
num_tokens = solara.reactive(10)
@solara.component
def Page():
  css = """
  .myclass{
    color:blue!important;
    font-size:2em;
  }
  """
  solara.Style(css)

  def generate(tokens):
    if seed.value:
      torch.manual_seed(seed_int.value)
      outputs = model.generate(tokens, do_sample=sample.value, max_new_tokens=num_tokens.value)
    else:
      outputs = model.generate(tokens, do_sample=sample.value, max_new_tokens=num_tokens.value)
    response = ""
    for output in outputs[0][len(tokens[0]):]:
      response += tokenizer.decode([output])
    return response

  with solara.Column(margin=10):
    title = "GPT-2"
    with solara.Head():
      solara.Title(title)
    solara.Markdown(f"#{title}")
    with solara.Row():
      checkbox_sample = solara.Checkbox(label="Sample", value=sample)
      checkbox_seed = solara.Checkbox(label="Seed", value=seed)
      if seed.value:
        solara.InputInt("Enter a seed value:", value=seed_int)
    solara.InputText("Enter text:", value=prompt, continuous_update=True)
    if prompt.value != "":
      tokens = tokenizer.encode(prompt.value, return_tensors="pt")
      def on_click():
        answer.value = ""
        answer.value = generate(tokens)
      with solara.Row():
        solara.Button(label="Generate Response", on_click=on_click)
        solara.SliderInt(f"number of new tokens: {num_tokens.value}", value=num_tokens, min=1, max=40)
      solara.Markdown("")
      with solara.Row():
        if answer.value != "":
          solara.Text(f"""{answer.value}""", classes=["myclass"])