benchmarks/speed/scripts/run_nlp.py
def _run_transformer_model(name: str, gpu) -> Callable[[List[str]], None]:
    """Run bare transformer model, outputting raw hidden-states"""
    from transformers import AutoTokenizer, AutoModel
    import torch
    if gpu:
        torch.device("cuda:0")

    tokenizer = AutoTokenizer.from_pretrained(name, use_fast=True)
    transformer = AutoModel.from_pretrained(name)
    if gpu:
        transformer = transformer.cuda()

    def run(texts: List[str], batch_size: int):
        transformer.eval()
        for batch in minibatch(texts, batch_size // 20):
            batch = tokenizer(batch, padding=True, truncation=True, return_tensors="pt")
            if gpu:
                batch["input_ids"] = batch["input_ids"].to("cuda:0")
                batch["attention_mask"] = batch["attention_mask"].to("cuda:0")
            transformer(**batch)

    return run
Run bare transformer model, outputting raw hidden-states
Run bare transformer model, outputting raw hidden-states
Run bare transformer model outputting raw hidden - states
def _run_transformer_model(name: str, gpu) -> Callable[[List[str]], None]:
    from transformers import AutoTokenizer, AutoModel
    import torch
    if gpu:
        torch.device("cuda:0")
    tokenizer = AutoTokenizer.from_pretrained(name, use_fast=True)
    transformer = AutoModel.from_pretrained(name)
    if gpu:
        transformer = transformer.cuda()
    def run(texts: List[str], batch_size: int):
        transformer.eval()
        for batch in minibatch(texts, batch_size // 20):
            batch = tokenizer(batch, padding=True, truncation=True, return_tensors="pt")
            if gpu:
                batch["input_ids"] = batch["input_ids"].to("cuda:0")
                batch["attention_mask"] = batch["attention_mask"].to("cuda:0")
            transformer(**batch)
    return run
def _run_transformer_model ( name : str , gpu ) -> Callable [ [ List [ str ] ] , None ] : from transformers import AutoTokenizer , AutoModel import torch if gpu : torch . device ( "cuda:0" ) tokenizer = AutoTokenizer . from_pretrained ( name , use_fast = True ) transformer = AutoModel . from_pretrained ( name ) if gpu : transformer = transformer . cuda ( ) def run ( texts : List [ str ] , batch_size : int ) : transformer . eval ( ) for batch in minibatch ( texts , batch_size // 20 ) : batch = tokenizer ( batch , padding = True , truncation = True , return_tensors = "pt" ) if gpu : batch [ "input_ids" ] = batch [ "input_ids" ] . to ( "cuda:0" ) batch [ "attention_mask" ] = batch [ "attention_mask" ] . to ( "cuda:0" ) transformer ( ** batch ) return run
Run bare transformer model, outputting raw hidden-states
Run bare transformer model outputting raw hidden - states
"""Run bare transformer model, outputting raw hidden-states"""
name
str
gpu