# -*- coding: utf-8 -*-
"""gradio_with_CodeGen.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1sZPTjB9cF90Ivu8LltJPEYSawXCE1LRv

# Interacting with [CodeGen](https://github.com/salesforce/CodeGen/)
"""



# Commented out IPython magic to ensure Python compatibility.
 
!git clone https://github.com/salesforce/CodeGen
# %cd CodeGen
!pip install --upgrade pip setuptools
!pip install gradio
!pip install -r requirements.txt

chosen_model = "codegen-350M-nl" #@param ["codegen-350M-nl", "codegen-350M-multi", "codegen-350M-mono", "codegen-2B-nl", "codegen-2B-multi", "codegen-2B-mono", "codegen-6B-nl", "codegen-6B-multi", "codegen-6B-mono", "codegen-16B-nl", "codegen-16B-multi", "codegen-16B-mono"]
fp16 = True #param {type:"boolean"}

import os

if not os.path.exists(f'./checkpoints/{chosen_model}'):
  !wget -P checkpoints https://storage.googleapis.com/sfr-codegen-research/checkpoints/{chosen_model}.tar.gz && tar -xvf checkpoints/{chosen_model}.tar.gz -C checkpoints/


import torch
from jaxformer.hf.sample import truncate as do_truncate
from jaxformer.hf.sample import set_env, set_seed, print_time, create_model, create_custom_gpt2_tokenizer, create_tokenizer, sample

# (0) constants

models_nl = ['codegen-350M-nl', 'codegen-2B-nl', 'codegen-6B-nl', 'codegen-16B-nl']
models_pl = ['codegen-350M-multi', 'codegen-2B-multi', 'codegen-6B-multi', 'codegen-16B-multi', 'codegen-350M-mono', 'codegen-2B-mono', 'codegen-6B-mono', 'codegen-16B-mono']
models = models_nl + models_pl


# (2) preamble

set_env()

pad = 50256
# device = torch.device('cuda:0')
device = torch.device("cpu")
ckpt = f'./checkpoints/{chosen_model}'

# if device.type == "cpu":
#   print()
#   print("force full precision for cpu!!")
#   print()
  
fp16 = False


# (3) load

with print_time('loading parameters'):
  model = create_model(ckpt=ckpt, fp16=fp16).to(device)


with print_time('loading tokenizer'):
  if chosen_model in models_pl:
    tokenizer = create_custom_gpt2_tokenizer()
  else:
    tokenizer = create_tokenizer()
  tokenizer.padding_side = 'left'
  tokenizer.pad_token = pad

def codegen(context):
 #param {type:"string"}
  
  rng_seed = 42 #param {type:"integer"}
  rng_deterministic = True #param {type:"boolean"}
  p = 0.95 #param {type:"number"}
  t = 0.1 #param {type:"number"}
  max_length = 128 #param {type:"integer"}
  batch_size = 1 #param {type:"integer"}
  set_seed(rng_seed, deterministic=rng_deterministic)

  # (4) sample

  with print_time('sampling'):
    completion = sample(device=device, model=model, tokenizer=tokenizer, context=context, pad_token_id=pad, num_return_sequences=batch_size, temp=t, top_p=p, max_length_sample=max_length)[0]
    truncation = do_truncate(completion)

    # print('=' * 100)
    # print(completion)
    # print('=' * 100)
    # print(context+truncation)
    # print('=' * 100)
      

  return completion

# !python -m jaxformer.hf.sample --model $chosen_model \
#                  --rng-seed $rng_seed \
#                  --p $p \
#                  --t $t \
#                  --max-length $max_length \
#                  --batch-size $batch_size \
#                  --context '$context'

# context = "def hello_world():"
# codegen(context)

import numpy as np
import gradio as gr
 

iface = gr.Interface(
    codegen,
    [   gr.inputs.Textbox(type='str',   label="input prompt"),
    ],
    "text",
)

iface.launch(debug=True)