近日，兩個外國小哥Kartik Godawat 和 Deepak Rawat 開發了一個 Jupyter 插件Text2Code，可以將自然語言查詢轉換成相關的 Python 數據分析代碼。

如果能創建一個桌面軟件，將自然語言直接轉換成相關的 Python 數據分析代碼，工作就方便了。

這不，有倆「好事」的程序員耐不住寂寞，把這個工具做出來。

靈感來自GPT-3，自然語言直接轉代碼

2020年6月，OpenAI 推出了GPT-3，它不僅具有未來的 NLP 相關的諸多功能，而且還能夠生成 React 代碼和shell命令。

倆人從中得到了靈感，他們意識到在做數據分析的時候，我們經常忘記不經常使用的 pandas 命令或者語法，需要去搜索，從 StackOverflow 復制代碼然后需要相應地修改變量和列名。

最初他們試圖把這個問題作為一個聊天機器人來解決，并試圖使用Rasa，但因為缺乏合適的訓練數據而夭折了。

他們決定開發一個監督學習模型，該模型可以吃進自然語言和代碼的對應語料，然后進行訓練，完整的pipeline包含了下面幾個步驟：

生成訓練數據

為了模擬終端的用戶想向系統查詢的內容，我們開始嘗試用英語描述一些命令的格式。

例如：

display a line plot showing $colname on y-axis and $colname on x-axis from $varname

然后，我們通過使用一個非常簡單的生成器來替換 $colname 和 $varname 來生成數據，以獲得訓練集中的變量。

意圖匹配

在生成數據之后，需要為特定的意圖映射到一個唯一的「意圖id」，并使用通用語句編碼器獲取用戶query的embedding，然后用我們預先定義好的意圖query(生成的數據)來得到余弦距離。

「通用句子編碼器」-Universal Sentence Encoder類似于 word2vec，會生成相應的embedding，唯一的區別是這里的嵌入是用于句子而不是單詞。

命名實體識別

相同的生成數據可以被用來訓練一個自定義的實體識別模型，這個模型可以用來檢測columns，varaibles，library的名字。

為此，作者還研究了 HuggingFace 模型，但最終決定使用 Spacy 來訓練模型，主要是因為 HuggingFace 模型是基于Transformer的模型，與 Spacy 相比有點過于復雜繁重。

填充模板

一旦實體被正確識別并且意圖被正確匹配，填充模板就變得非常容易。

例如，“ show 5 rows from df”這個查詢將生成兩個實體: 一個變量和一個數值。這個模板代碼編寫起來很簡單。

與Jupyter結合

這是所有步驟里最復雜的一步，因為為 Jupyter 編寫如此復雜的擴展有點棘手，而且幾乎沒有文檔或示例參考。

他們通過一些嘗試并參考已經存在的擴展，最終將所有內容包裝成一個單獨的 Python 包，并且可以直接通過 pip 來安裝。

單單一個Python包并不能直接使用，于是他們創建了一個前端以及一個服務器擴展，當 jupyter notebook啟動時，就會被加載。前端向服務器發送查詢以獲取生成的模板代碼，然后將其插入notebook的對應單元并執行它。

Text2Code的演示

模型也會失效，但數據分析師真的省事了

就像許多機器學習模型一樣，有時候意圖匹配和命名實體識別效果會很差，即使這個意圖對人來說非常簡單。

有時也會識別不到意圖，就無法生成正確的代碼，作者還考慮使用下面的方法來進一步改進插件的效果。

收集/生成高質量的英語訓練數據，可以考慮從quroa，StackOverflow爬取更多的高贊回答，嘗試用不同的方式來描述相同的內容，增強數據；收集真實世界的變量名和庫名，而不是隨機生成，使用基于Transformer的模型嘗試命名實體識別。

如果這個模型訓練的足夠好，對數據分析師來，能省不少事了。

項目開源地址：

https://github.com/deepklarity/jupyter-text2code