Roberta Go

🚀 roberta-go

這是一個基於 CodeSearchNet 數據集 針對 Golang 掩碼語言模型任務進行預訓練的 roberta 版本，可用於填充 Go 代碼中的掩碼詞。

🚀 快速開始

安裝必要的包

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

💻 使用示例

基礎用法

code = """
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()

pred = {x["token_str"].replace("Ġ", ""): x["score"] for x in fill_mask(code)}
sorted(pred.items(), key=lambda kv: kv[1], reverse=True)
[('GOOS', 0.11810332536697388),
 ('FileInfo', 0.04276798665523529),
 ('Stdout', 0.03572738170623779),
 ('Getenv', 0.025064032524824142),
 ('FileMode', 0.01462600938975811)]

使用上述代碼，你可以使用該模型填充 Go 代碼中的掩碼詞。

📚 詳細文檔

整個訓練過程和超參數設置可在我的 GitHub 倉庫 中查看。

本項目由 Dor Bernsohn 創建。

📦 模型信息