Story : Plots , UX for LLM support in Translation [1]

[SoyGema]

Submission Type

• Feature Request
• Discussion

Context

Neural Machine Translation Experiment Tracking scenario. Repo
ML Area : Transfer learning with LLMs. Fine-tunning of t5-small with opus100 dataset from HF
Use DVC VS Code extension with dvclive experiment tracking scenario

tensorflow-macos==2.9.0
tensorflow-metal==0.5.0
transformers==4.32.0.dev0
dvc==3.x.x
dvclive==TBD

Impact

• Explainability of Transformer models with visualizations ( Attention Heads )
• Plot text and translated text ( data string format) to evaluate empirically how well LLMs perform (scalable to another use cases)
• Teach results at PyCon Spain Keynote 23
• Comparison of two ML experiment tracking frameworks: mlflow VS dvc from DS perspective

Issue creator Goal

• Discuss implementation
• Offer support in other LLM use cases
• Get help / possible template / contrib opportunity

1. Description

Attention Heads

What to plot and why ?

Attention Heads. It allows us to see how the words are mapping with respect to two different languages.
At a high level from the Machine Learning perspective it allows us to analyze how well verbs are translated, how well the model is understanding prepositions in one language to respect to another, etc... This visualization might be shown after fine-tunning.

Code snippets useful for visualization

Main plot function . Previous Script

def plot_attention_head(in_tokens, translated_tokens, attention):
  # The model didn't generate `<START>` in the output. Skip it.
  translated_tokens = translated_tokens[1:]

  ax = plt.gca()
  ax.matshow(attention)
  ax.set_xticks(range(len(in_tokens)))
  ax.set_yticks(range(len(translated_tokens)))

  labels = [label.decode('utf-8') for label in in_tokens.numpy()]
  ax.set_xticklabels(
      labels, rotation=90)

  labels = [label.decode('utf-8') for label in translated_tokens.numpy()]
  ax.set_yticklabels(labels)

Questions :
Would it be possible to plot this in the plot section in the extension easily ?
From the DS standpoint, Vega is not the standard. In case of develop it by myself, any idea of how long do you estimate this would take ?
DO you know if something more interactive like bertviz is possible?

2. Description

Text

What to plot and why ?

The Goal is to plot the sentence and the translated sentence in the plots section in VSCode with the extension, but Im assuming that this would entail another template. In translation, the sentence in one language (input) and another language (output once the model is fine-tunned) , to evaluate empirically how well the model is translating a set of sentences that might be relevant.

Questions :
Can I log it via dvclive as well?

Thanks in advance!