layoutlmv2

PHOTO EMBED

Thu Apr 13 2023 08:14:53 GMT+0000 (Coordinated Universal Time)

Saved by @mehla99_shubham #python

from layoutlm import FunsdDataset, LayoutlmConfig, LayoutlmForTokenClassification 
from transformers import BertTokenizer,AdamW 
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler 
from torch.utils.data.distributed import DistributedSampler 
import torch 
from tqdm import tqdm, trange

MODEL_CLASSES = { "layoutlm": (LayoutlmConfig, LayoutlmForTokenClassification, BertTokenizer), }

def train( train_dataset, model, tokenizer, labels, pad_token_label_id):

    """ Train the model """
    if torch.cuda.is_available():
    	device = torch.device("cuda") 
		print("GPU is available")
    else:
    	device = torch.device("cpu")
	    print("GPU is not available, using CPU instead")

    train_sampler = RandomSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, 		   batch_size=args.train_batch_size, collate_fn=None )
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer = AdamW( lr=learning_rate, eps=adam_epsilon)
    model = torch.nn.Module(model, find_unused_parameters=True)

    global_step, tr_loss = 0, 0.0
    model.zero_grad()
    train_iterator = trange(num_train_epochs, desc="Epoch")
    for in trainiterator:
    	epoch_iterator = tqdm(train_dataloader, desc="Iteration")
    	for step, batch in enumerate(epoch_iterator):
    		model.train()
    		inputs = {"input_ids": batch[0].to(device), "attention_mask": batch[1].to(device), 	"labels": batch[3].to(device)}
    		inputs["bbox"] = batch[4].to(device)
    		inputs["token_type_ids"] = batch[2].to(device)

    outputs = model(**inputs)
    loss = outputs[0]
    loss.backward()
    tr_loss += loss.item()

    optimizer.step() # scheduler.step() # Update learning rate schedule
    model.zero_grad()
    global_step += 1

    return global_step, tr_loss / global_step

def main():

    if torch.cuda.is_available():
    	device = torch.device("cuda")
    	print("GPU is available")
    else:
        device = torch.device("cpu")
        print("GPU is not available, using CPU instead")

    labels = get_labels(labels) # in our case labels will be x-axis,y-axis,title
    num_labels = len(labels) # Use cross entropy ignore index as padding label id so that only real 	label ids contribute to the loss later
    pad_token_label_id = CrossEntropyLoss().ignore_index

    config = config_class.from_pretrained( "layoutlm-base-uncased/", num_labels=num_labels, 		force_download = True, ignore_mismatched_sizes=True, cache_dir= cache_dir_path else None, ) 	
	tokenizer = tokenizer_class.from_pretrained( "microsoft/layoutlm-base-uncased", 	 do_lower_case=True, force_download = True, ignore_mismatched_sizes=True, cache_dir= cache_dir_path else None, )

    model = model_class.from_pretrained( "layoutlm-base-uncased/", config=config, )
    model.to(args.device)

    train_dataset = FunsdDataset( args, tokenizer, labels, pad_token_label_id, mode="train" )
    global_step, tr_loss = train( args, train_dataset, model, tokenizer, labels, pad_token_label_id )

    tokenizer = tokenizer_class.from_pretrained( "microsoft/layoutlm-base-uncased",force_download = True, do_lower_case=args.do_lower_case,ignore_mismatched_sizes=True)

    model = model_class.from_pretrained(args.output_dir)
    model.to(args.device)
    result, predictions = evaluate( args, model, tokenizer, labels, pad_token_label_id, mode="test" )

    return result,predictions
content_copyCOPY