Upload gradio_with_codegen.py

gradio_with_codegen.py

@@ -0,0 +1,121 @@
+# -*- 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)