Semantic Code Retrieval using Transformers
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.




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).