Semantic Code Retrieval using Transformers
Author: celsofranssa
Pull request: no
Labels: name: wontfix

I am entering the world of transformers and would like to use some architectures to create a semantic search engine to retrieve source code (Python, Javascript, Ruby, Go, Java, and PHP code).

Currently, the [dataset ](https://github.com/github/CodeSearchNet#data-details)contains 2 million pairs **(code, docstring)**, where code is a list of tokens from a method or function and docstring is a short description of the code in natural language.

As a starting point, it would be interesting to construct a model architecture that receives the code and the docstring **([ [code], [docstring] ])** as input example and outputs the code embedding and docstring embedding. Using cosine similarity as loss function the model could be fine-tuned to encode both code and docstring to the same embedding space. As shown in the figure below:

<pre>            <img src="https://i.stack.imgur.com/4fx3h.png" width="480">
</pre>

I started reading and tokenizing the dataset:

```python
    from transformers import BertTokenizer
    # reads a list of [[code], [docstring]]
    reader = CodeDocstringReader(dataset_path)
    
    # loads tokenizer
    model_name = "bert-base-uncased"
    tokenizer = BertTokenizer.from_pretrained(model_name, do_lower_case=True)

    # returns a list of tokenized examples
    # [[code_tokes_ids], [docstring_tokens_ids]]
    tokenized_features = tokenizer_examples(
        reader.get_examples(),
        tokenizer
    )
```

The definition and training of the model are still incomplete, but it is outlined as:

```python
import tensorflow as tf
from transformers import BertModel


class JointEncoder(tf.keras.Model):
    """Encodes the code and docstring into an same space of embeddings."""

    def __init__(self,
                 path,
                 name="jointencoder"):
        super(JointEncoder, self).__init__(name=name)
        self.bert = BertModel.from_pretrained(path)

    def call(self, inputs):
        """Returns code and docstring embeddings"""
         ...
         code_embedding = ..
         docstring_embedding = ..

        return code_embedding, docstring_embedding 

```


However, I'm stuck on how to code this simple architecture. Could you give me some directions? 

Thanks in advance.





Comments:
You might be interested in:
- https://huggingface.co/huggingface/CodeBERTa-small-v1#codeberta and https://huggingface.co/huggingface/CodeBERTa-language-id
- more generally, https://huggingface.co/blog/how-to-train

Wow, it helped a lot @julien-c. Basically, did you train a language model using CodeSearchNet dataset?

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


@julien-c,

Is there an approach where this tutorial https://huggingface.co/blog/how-to-train is trained over TPU cores?

maybe @LysandreJik or @sgugger have a link to a notebook?

I haven't tried training in notebooks on TPU, only with the example scripts.

Unfortunately I only have notebooks that run the example scripts on TPU, nothing similar to the `how-to-train` blogpost.

Thanks @julien-c, @sgugger and @LysandreJik,

Maybe I can adapt the [Language Modeling example script](https://github.com/huggingface/transformers/tree/master/examples/language-modeling) applying [Pytorch Lightning](https://www.pytorchlightning.ai/) approach which easily support TPUs.

If you're using the language modeling script, running it on TPUs is supported, just follow the instructions [here](https://github.com/huggingface/transformers/tree/master/examples#running-on-tpus).

> If you're using the language modeling script, running it on TPUs is supported, just follow the instructions [here](https://github.com/huggingface/transformers/tree/master/examples#running-on-tpus).

Great, but it runs on Google Colab?

Yes it does run on colab!

You can find an example running the `run_glue.py` script [here](https://colab.research.google.com/drive/15q6UUzwNugWvVXNfOkWlGCaKIUvLZpxd?usp=sharing). You can do the same with the language modeling script! (Cloning the repository and running the script from there would be cleaner than `wget`ting all the files like it's done in this colab, though)

> You can find an example running the `run_glue.py` script [here](https://colab.research.google.com/drive/15q6UUzwNugWvVXNfOkWlGCaKIUvLZpxd?usp=sharing). You can do the same with the language modeling script! (Cloning the repository and running the script from there would be cleaner than `wget`ting all the files like it's done in this colab, though)

Fantastic! That was of great help.

@LysandreJik, @sgugger,

Unfortunately, even using a small dataset (~400MB), the Colab killed the process due to the use of all available RAM (12.72GB).