Skip to content

Commit

Permalink
added hugging face example (#81)
Browse files Browse the repository at this point in the history
Co-authored-by: John Calderon <john.calderon@centml.ai>
  • Loading branch information
johncalesp and John Calderon committed Mar 5, 2024
1 parent f371351 commit 271a2b3
Showing 1 changed file with 55 additions and 0 deletions.
55 changes: 55 additions & 0 deletions examples/huggingface/entry_point.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,55 @@
from transformers import (
get_linear_schedule_with_warmup,
AutoModelForCausalLM,
Trainer,
)
import torch
import torch.optim as optim

model_id = "roberta-base"


def deepview_model_provider():
return AutoModelForCausalLM.from_pretrained(model_id, is_decoder=True).cuda()


def deepview_input_provider(batch_size=2):
vocab_size = 30522
src_seq_len = 512
tgt_seq_len = 512

device = torch.device("cuda")

source = torch.randint(
low=0,
high=vocab_size,
size=(batch_size, src_seq_len),
dtype=torch.int64,
device=device,
)
target = torch.randint(
low=0,
high=vocab_size,
size=(batch_size, tgt_seq_len),
dtype=torch.int64,
device=device,
)
return (source, target)


def deepview_iteration_provider(model):
model.parameters()
optimizer = optim.AdamW(
params=model.parameters(),
betas=(0.9, 0.999),
eps=1e-6,
weight_decay=0.01,
lr=1e-4,
)
scheduler = get_linear_schedule_with_warmup(optimizer, 10000, 500000)
trainer = Trainer(model=model, optimizers=(optimizer, scheduler))

def iteration(source, label):
trainer.training_step(model, {"input_ids": source, "labels": label})

return iteration

0 comments on commit 271a2b3

Please sign in to comment.