Hugging Face's logo

Spaces:
Xhaheen
/
codegen-350M-nl
No application file

App Files Community
codegen-350M-nl
File size: 3,455 Bytes
e8bf8e8
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
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)