Support for COMET models

Model Support
1

I am trying to use NNsight for interpretability of machine translation evaluation metrics. Specifically, I am interested in the COMET metric: GitHub - Unbabel/COMET: A Neural Framework for MT Evaluation

The metric has its own model loading and inference functions and it doesn’t use the huggingface Transformers standard functions.

Where should I go from here?

2

Hi @Wafaa , NNsight supports any PyTorch model in theory! Looking at the comet code id try something like this:

from comet import download_model, load_from_checkpoint

model_path = download_model("Unbabel/wmt22-comet-da")
model = load_from_checkpoint(model_path)
data = [
    {
        "src": "Dem Feuer konnte Einhalt geboten werden",
        "mt": "The fire could be stopped",
        "ref": "They were able to control the fire."
    },
    {
        "src": "Schulen und Kindergärten wurden eröffnet.",
        "mt": "Schools and kindergartens were open",
        "ref": "Schools and kindergartens opened"
    }
]

from nnsight import NNsight

model = NNsight(model)

with model.predict(data, batch_size=8, gpus=1) as tracer:
    model_output = tracer.result.save()
print (model_output)

Please report back if this works!

1 Like
3

It works! thank you so much!

1 Like