Browse Source
Add graphcodebert
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
2 changed files with
9 additions and
3 deletions
-
README.md
-
codebert.py
|
|
|
@ -8,7 +8,8 @@ |
|
|
|
|
|
|
|
A code search operator takes a text string of programming language or natural language as an input |
|
|
|
and returns an embedding vector in ndarray which captures the input's core semantic elements. |
|
|
|
This operator is implemented with pre-trained models from [Huggingface Transformers](https://huggingface.co/docs/transformers). |
|
|
|
This operator is implemented with pre-trained [CodeBert](https://arxiv.org/pdf/2002.08155.pdf) |
|
|
|
or [GraphCodeBert](https://arxiv.org/abs/2009.08366) models from [Huggingface Transformers](https://huggingface.co/docs/transformers). |
|
|
|
|
|
|
|
<br /> |
|
|
|
|
|
|
|
@ -64,7 +65,11 @@ The device to run model inference. |
|
|
|
The default value is None, which enables GPU if cuda is available. |
|
|
|
|
|
|
|
Supported model names: |
|
|
|
|
|
|
|
- huggingface/CodeBERTa-small-v1 |
|
|
|
- microsoft/codebert-base |
|
|
|
- microsoft/codebert-base-mlm |
|
|
|
- mrm8488/codebert-base-finetuned-stackoverflow-ner |
|
|
|
- microsoft/graphcodebert-base |
|
|
|
|
|
|
|
<br /> |
|
|
|
|
|
|
|
|
|
|
|
@ -143,7 +143,8 @@ class CodeBert(NNOperator): |
|
|
|
'huggingface/CodeBERTa-small-v1', |
|
|
|
'microsoft/codebert-base', |
|
|
|
'microsoft/codebert-base-mlm', |
|
|
|
'mrm8488/codebert-base-finetuned-stackoverflow-ner' |
|
|
|
'mrm8488/codebert-base-finetuned-stackoverflow-ner', |
|
|
|
'microsoft/graphcodebert-base' |
|
|
|
] |
|
|
|
full_list.sort() |
|
|
|
if format is None: |
|
|
|
|
