Fix GPU memory leak in TextPairRegressor

[MattGPT-ai]

#3487

This addresses the above memory leak by saving a reference to the TextPair when the concatenated Sentence is created. Its embeddings are explicitly cleared when clear_embeddings is called.

This also adds type hints to files relevant to this change.

[MattGPT-ai]

Just pushed a type hint change for a file that I didn't really touch, but I think is checking because I added type hints to Sentence:
predicted: List[List[Union[int, float]]] = [[] for _ in range(number_tokens)]

Not sure if it's the case that this should just be List[List[int]] and the line predicted[i].append(input_indices[i].item()) should just cast int(input_indices[i].item()) - but that seemed riskier since it's changing code and not just the type hint.
Any comments here?

[alanakbik]

Thanks for adding this! Mostly looks good. Just please correct the return signature of the embedding property, as indicated in the comment above.

I also added a suggestion on checking for the concatenated data point, which may give a minimal speed-up during training (though likely negligible so up to you if you want to include it).

Regarding the lemmatizer, I'm not so sure. Probably best to err on the side of caution and use the old, underspecified signature.

[mattb-zip]

I have addressed all of the requests I believe, let me know if you can approve or request any further changes

[MattGPT-ai]

I just realized that this same issue exists for TextPairClassifier

The simple fix is to just apply the same logic to _get_embedding_for_data_point -- however, since these two functions are identical, it might be good to do a small refactor where they each call the same common function for embedding a TextPair.

Since these inherit from different base classes, it doesn't seem like the best idea to try to do this via inheritance

[alanakbik]

@MattGPT-ai thanks for fixing this and adding the type hints!

I'll merge this now. The same bug in TextPairClassifier could be fixed similarly. Agree that addressing this via inheritance would be overkill since this only affects two classes (and doing too much via inheritance tends to make code less intuitive).

[MattGPT-ai]

I actually pushed an additional commit to fix it in TextPairClassifier, just did the same fix. I have an idea for refactoring this: make _get_embedding_for_data_point a function in TextPair - right now it's just a type definition, but it could be made a class that inherits DataPair[Sentence, Sentence]. Then the models that use TextPair would just use its own function to embed. Could have a single function or two functions for embedding separately or jointly.

This would give better control over how these embeddings are set in the TextPair. One problem I thought of is that if we're saving the concatenated sentence, you could have an inconsistency if you make a change to one of the texts in the pair, and it would not be updated accordingly. This refactor would allow us to have a setter function for each member of the pair that would automatically update the concatenated sentence. I think right now we're relying too much on the models to handle all of this, but all this functionality should exist in the data class independent of the models.

Currently in DataPair, you have the property embedding that returns the concatenation of the two member embeddings concatenated, which is more in line with doing separate embeddings. That might be a decent default, but I think it makes sense to override that for the case of TextPair, and it might even be good to have a property that sets if its saved embedding is separate or joint. Not sure about that, but your thoughts would be much appreciated. If you think that makes sense, I could push a PR and maybe create a separate issue

[alanakbik]

Hi @MattGPT-ai, is it a good idea and agree that saving the concatenated data point may lead to problematic inconsistencies.

I'm a bit worried though that moving too much logic into the data class will cause the logic to become too distributed. Right now, the TextPairClassifier makes all model-relevant decisions through _get_embedding_for_data_point: whether or not to concat, what separator to use, etc. Perhaps you could prepare a PR and we can discuss it?