microsoft/unixcoder-base

1---2language:3- en4license: apache-2.05---6 7# Model Card for UniXcoder-base8 9 10 11# Model Details12 13## Model Description14UniXcoder is a unified cross-modal pre-trained model that leverages multimodal data (i.e. code comment and AST) to pretrain code representation. 15 16- **Developed by:** Microsoft Team 17- **Shared by [Optional]:** Hugging Face18- **Model type:** Feature Engineering19- **Language(s) (NLP):** en20- **License:** Apache-2.021- **Related Models:**22  - **Parent Model:** RoBERTa23- **Resources for more information:**24    - [Associated Paper](https://arxiv.org/abs/2203.03850)25 26# Uses27 28## 1. Dependency29 30- pip install torch31- pip install transformers32 33## 2. Quick Tour34We implement a class to use UniXcoder and you can follow the code to build UniXcoder.35You can download the class by36```shell37wget https://raw.githubusercontent.com/microsoft/CodeBERT/master/UniXcoder/unixcoder.py38```39 40```python41import torch42from unixcoder import UniXcoder43 44device = torch.device("cuda" if torch.cuda.is_available() else "cpu")45model = UniXcoder("microsoft/unixcoder-base")46model.to(device)47```48 49In the following, we will give zero-shot examples for several tasks under different mode, including **code search (encoder-only)**, **code completion (decoder-only)**, **function name prediction (encoder-decoder)** , **API recommendation (encoder-decoder)**, **code summarization (encoder-decoder)**.50 51## 3. Encoder-only Mode52 53For encoder-only mode, we give an example of **code search**.54 55### 1) Code and NL Embeddings56 57Here, we give an example to obtain code fragment embedding from CodeBERT.58 59```python60# Encode maximum function61func = "def f(a,b): if a>b: return a else return b"62tokens_ids = model.tokenize([func],max_length=512,mode="<encoder-only>")63source_ids = torch.tensor(tokens_ids).to(device)64tokens_embeddings,max_func_embedding = model(source_ids)65 66# Encode minimum function67func = "def f(a,b): if a<b: return a else return b"68tokens_ids = model.tokenize([func],max_length=512,mode="<encoder-only>")69source_ids = torch.tensor(tokens_ids).to(device)70tokens_embeddings,min_func_embedding = model(source_ids)71 72# Encode NL73nl = "return maximum value"74tokens_ids = model.tokenize([nl],max_length=512,mode="<encoder-only>")75source_ids = torch.tensor(tokens_ids).to(device)76tokens_embeddings,nl_embedding = model(source_ids)77 78print(max_func_embedding.shape)79print(max_func_embedding)80```81 82```python83torch.Size([1, 768])84tensor([[ 8.6533e-01, -1.9796e+00, -8.6849e-01,  4.2652e-01, -5.3696e-01,85         -1.5521e-01,  5.3770e-01,  3.4199e-01,  3.6305e-01, -3.9391e-01,86         -1.1816e+00,  2.6010e+00, -7.7133e-01,  1.8441e+00,  2.3645e+00,87				 ...,88         -2.9188e+00,  1.2555e+00, -1.9953e+00, -1.9795e+00,  1.7279e+00,89          6.4590e-01, -5.2769e-02,  2.4965e-01,  2.3962e-02,  5.9996e-02,90          2.5659e+00,  3.6533e+00,  2.0301e+00]], device='cuda:0',91       grad_fn=<DivBackward0>)92```93 94### 2) Similarity between code and NL95 96Now, we calculate cosine similarity between NL and two functions. Although the difference of two functions is only a operator (```<``` and ```>```), UniXcoder can distinguish them.97 98```python99# Normalize embedding100norm_max_func_embedding = torch.nn.functional.normalize(max_func_embedding, p=2, dim=1)101norm_min_func_embedding = torch.nn.functional.normalize(min_func_embedding, p=2, dim=1)102norm_nl_embedding = torch.nn.functional.normalize(nl_embedding, p=2, dim=1)103 104max_func_nl_similarity = torch.einsum("ac,bc->ab",norm_max_func_embedding,norm_nl_embedding)105min_func_nl_similarity = torch.einsum("ac,bc->ab",norm_min_func_embedding,norm_nl_embedding)106 107print(max_func_nl_similarity)108print(min_func_nl_similarity)109```110 111```python112tensor([[0.3002]], device='cuda:0', grad_fn=<ViewBackward>)113tensor([[0.1881]], device='cuda:0', grad_fn=<ViewBackward>)114```115 116## 3. Decoder-only Mode117 118For decoder-only mode, we give an example of **code completion**.119 120```python121context = """122def f(data,file_path):123    # write json data into file_path in python language124"""125tokens_ids = model.tokenize([context],max_length=512,mode="<decoder-only>")126source_ids = torch.tensor(tokens_ids).to(device)127prediction_ids = model.generate(source_ids, decoder_only=True, beam_size=3, max_length=128)128predictions = model.decode(prediction_ids)129print(context+predictions[0][0])130```131 132```python133def f(data,file_path):134    # write json data into file_path in python language135    data = json.dumps(data)136    with open(file_path, 'w') as f:137        f.write(data)138```139 140## 4. Encoder-Decoder Mode141 142For encoder-decoder mode, we give two examples including: **function name prediction**, **API recommendation**, **code summarization**.143 144### 1) **Function Name Prediction**145 146```python147context = """148def <mask0>(data,file_path):149    data = json.dumps(data)150    with open(file_path, 'w') as f:151        f.write(data)152"""153tokens_ids = model.tokenize([context],max_length=512,mode="<encoder-decoder>")154source_ids = torch.tensor(tokens_ids).to(device)155prediction_ids = model.generate(source_ids, decoder_only=False, beam_size=3, max_length=128)156predictions = model.decode(prediction_ids)157print([x.replace("<mask0>","").strip() for x in predictions[0]])158```159 160```python161['write_json', 'write_file', 'to_json']162```163 164### 2) API Recommendation165 166```python167context = """168def write_json(data,file_path):169    data = <mask0>(data)170    with open(file_path, 'w') as f:171        f.write(data)172"""173tokens_ids = model.tokenize([context],max_length=512,mode="<encoder-decoder>")174source_ids = torch.tensor(tokens_ids).to(device)175prediction_ids = model.generate(source_ids, decoder_only=False, beam_size=3, max_length=128)176predictions = model.decode(prediction_ids)177print([x.replace("<mask0>","").strip() for x in predictions[0]])178```179 180```python181['json.dumps', 'json.loads', 'str']182```183 184### 3) Code Summarization185 186```python187context = """188# <mask0>189def write_json(data,file_path):190    data = json.dumps(data)191    with open(file_path, 'w') as f:192        f.write(data)193"""194tokens_ids = model.tokenize([context],max_length=512,mode="<encoder-decoder>")195source_ids = torch.tensor(tokens_ids).to(device)196prediction_ids = model.generate(source_ids, decoder_only=False, beam_size=3, max_length=128)197predictions = model.decode(prediction_ids)198print([x.replace("<mask0>","").strip() for x in predictions[0]])199```200 201```python202['Write JSON to file', 'Write json to file', 'Write a json file']203```204 205 206 207 208# Reference209If you use this code or UniXcoder, please consider citing us.210 211<pre><code>@article{guo2022unixcoder,212  title={UniXcoder: Unified Cross-Modal Pre-training for Code Representation},213  author={Guo, Daya and Lu, Shuai and Duan, Nan and Wang, Yanlin and Zhou, Ming and Yin, Jian},214  journal={arXiv preprint arXiv:2203.03850},215  year={2022}216}</code></pre>217 218 219 220