roberta-go

语言: Go 数据集:

• code_search_net

这是一个基于CodeSearchNet数据集预训练的roberta模型，专为Golang掩码语言模型任务设计。

加载模型的方法： （需要安装的包：!pip install transformers sentencepiece）

from transformers import AutoTokenizer, AutoModelWithLMHead, pipeline
tokenizer = AutoTokenizer.from_pretrained("dbernsohn/roberta-go")
model = AutoModelWithLMHead.from_pretrained("dbernsohn/roberta-go")

fill_mask = pipeline(
    "fill-mask",
    model=model,
    tokenizer=tokenizer
)

随后可以使用该模型来填充Go代码中的掩码词汇。

代码示例 = """
package main

import (
	"fmt"
	"runtime"
)

func main() {
	fmt.Print("Go runs on ")
	switch os := runtime.<mask>; os {
	case "darwin":
		fmt.Println("OS X.")
	case "linux":
		fmt.Println("Linux.")
	default:
		// freebsd, openbsd,
		// plan9, windows...
		fmt.Printf("%s.\n", os)
	}
}
""".lstrip()

预测结果 = {x["token_str"].replace("Ġ", ""): x["score"] for x in fill_mask(代码示例)}
排序结果 = sorted(预测结果.items(), key=lambda kv: kv[1], reverse=True)
[('GOOS', 0.11810332536697388),
 ('FileInfo', 0.04276798665523529),
 ('Stdout', 0.03572738170623779),
 ('Getenv', 0.025064032524824142),
 ('FileMode', 0.01462600938975811)]

完整的训练过程和超参数设置可在我的GitHub仓库中找到。

由Dor Bernsohn创建