validate_export.py

import torch
from transformers import AutoTokenizer

export_path = "examples/exports/distilbert-base-uncased_itc_quantized.pt"
model_name = "distilbert-base-uncased"

def main():
    print(f"Loading tokenizer for {model_name}")
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    text = "This is a validation sentence for the exported model."
    inputs = tokenizer(text, return_tensors="pt")
    input_ids = inputs['input_ids']
    attention_mask = inputs.get('attention_mask', None)

    print(f"Loading TorchScript model from {export_path}")
    m = torch.jit.load(export_path)
    m.eval()

    with torch.no_grad():
        try:
            if attention_mask is not None:
                out = m(input_ids, attention_mask)
            else:
                out = m(input_ids)
        except TypeError:
            # Some traced models expect a single arg
            out = m(input_ids)

    print("Output type:", type(out))
    try:
        print("Output shape:", out.shape)
        # print first 5 logits
        print("Sample logits (first 5):", out[0, :5].detach().cpu().numpy())
    except Exception as e:
        print("Could not print output shape/logits:", e)

if __name__ == '__main__':
    main()