Commit 5b1beceb authored by jon-tow's avatar jon-tow
Browse files

Move `TODO`s into one source location

parent ff83e9b2
...@@ -33,18 +33,19 @@ class HFLM(BaseLM): ...@@ -33,18 +33,19 @@ class HFLM(BaseLM):
else torch.device("cpu") else torch.device("cpu")
) )
# TODO: update this to be less of a hack once subfolder is fixed in HF
revision = revision + ("/" + subfolder if subfolder is not None else "")
self.gpt2 = transformers.AutoModelForCausalLM.from_pretrained( self.gpt2 = transformers.AutoModelForCausalLM.from_pretrained(
pretrained, pretrained,
# TODO: update this to be less of a hack once subfolder is fixed in HF revision=revision,
revision=revision + ("/" + subfolder if subfolder is not None else ""),
).to(self.device) ).to(self.device)
self.gpt2.eval() self.gpt2.eval()
# pretrained tokenizer for neo is broken for now so just hard-coding this to gpt2 # pretrained tokenizer for neo is broken for now so just hard-coding this to gpt2
self.tokenizer = transformers.AutoTokenizer.from_pretrained( self.tokenizer = transformers.AutoTokenizer.from_pretrained(
pretrained if tokenizer is None else tokenizer, pretrained if tokenizer is None else tokenizer,
# TODO: update this to be less of a hack once subfolder is fixed in HF revision=revision,
revision=revision + ("/" + subfolder if subfolder is not None else ""),
) )
assert isinstance( assert isinstance(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment