test.py

from transformers import pipeline                                                    
import time                                                                          
                                                                                     
n = 50                                                                               
nlp_token_class_cpu = pipeline("ner")                                                
                                                                                     
start = time.time()                                                                  
resp = nlp_token_class_cpu(["Hugging Face is a French company based in New-York."] * n)     
print("Device: CPU")                                                                 
print(f"No. examples: {n}")                                                          
print(f"Time taken: {time.time() - start}")                                              
                                                                                      
start = time.time()                                                                  
resp = nlp_token_class_cpu(["Hugging Face is a French company based in New-York."] * n, model_batch_size=2)     
print("Device: CPU (batched)")                                                       
print(f"No. examples: {n}")                                                          
print(f"Time taken: {time.time() - start}")                                          
                                                                                     
start = time.time()                                                                  
resp = nlp_token_class_cpu(                                                          
    ["Hugging Face is a French company based in New-York."] * (n - 1)                
    + [                                                                                
        "As he crossed toward the pharmacy at the corner he involuntarily turned his head because of a burst of light that had ricocheted from his temple, and saw, with that quick smile with which we greet a rainbow or a rose"    
    ],                                                                               
    model_batch_size=2,                                                              
)                                                                                    
print("Device: CPU (batched 2nd)")                                                   
print(f"No. examples: {n}")                                                          
print(f"Time taken: {time.time() - start}")                                          
                                                                                     
                                                                                     
print("-" * 50)                                                                      
                                                                                                                              
                                                                                                                              
nlp_token_class_gpu = pipeline("ner", device=0)                                                                               
start = time.time()                                                                                                           
resp = nlp_token_class_gpu(["Hugging Face is a French company based in New-York."] * n)                                       
print("Device: GPU")                                                                                                          
print(f"No. examples: {n}")                                                                                                   
print(f"Time taken: {time.time() - start}")                                          
start = time.time()                                                                                                           
                                                                                                                              
resp = nlp_token_class_gpu(["Hugging Face is a French company based in New-York."] * n, model_batch_size=2)                   
print("Device: GPU (batched)")                                                                                                
print(f"No. examples: {n}")                                                                                                   
print(f"Time taken: {time.time() - start}")                                                                                   
                                                                                                                              
start = time.time()                                                                                                           
resp = nlp_token_class_gpu(                                                                                                   
    ["Hugging Face is a French company based in New-York."] * (n - 1)                                                         
    + [                                                                              
        "As he crossed toward the pharmacy at the corner he involuntarily turned his head because of a burst of light that had ricocheted from his temple, and saw, with that quick smile with which we greet a rainbow or a rose, a blindingly white parallelogram of sky being unloaded from the van—a dresser with mirrors across which, as across a cinema screen, passed a flawlessly clear reflection of boughs sliding and swaying not arboreally, but with a human vacillation, produced by the nature of those who were carrying this sky, these boughs, this gliding façade."                                                               
    ],                                                                                                                        
    model_batch_size=2,                                                                                                       
)                                                                                                                             
print("Device: GPU (batched 2nd)")                                                                                            
print(f"No. examples: {n}")                                                                                                   
print(f"Time taken: {time.time() - start}")