jupytercoder/src/api.js at 9d0f613f0f4645ad9bae03b4ab76c1f0e2cea680 · bigcode-project/jupytercoder · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
const api = {
    sendToOpenAI(prompt, apiKey, modelType) { },
    sendToBigcode(code, url, token) { }
};


const cleanUpBigcodeOutput = (suggestion, isLastLine) => {
    suggestion = suggestion.replace("# -*- coding: utf-8 -*-\n\n", "").replace(/\t/g, '    ')
    let outPutIndex = suggestion.indexOf('<jupyter_output>')

    if (outPutIndex != -1){
        suggestion = suggestion.slice(0, outPutIndex)
    }

    const unnecessaryTag = '<|endoftext|>'
    const unnecessaryTagIndex = suggestion.indexOf(unnecessaryTag)


    suggestion = unnecessaryTagIndex == -1 ? suggestion : suggestion.slice(0, unnecessaryTagIndex)

    return isLastLine ? suggestion : suggestion.replace("\n","")
}


const cleanUpOpenaiOutput = (suggestion, isLastLine) => {
    if (!isLastLine){
        return suggestion.split("\n").length ==  1 ? suggestion : suggestion.split("\n")[0]
    }

    suggestion = suggestion.replace(/\u200B/g, '')
    let outPutIndex = suggestion.indexOf("\n\n§ Output")


    if (outPutIndex == -1) {
        let outPutIndex = suggestion.indexOf("\n\n# Output")

        return outPutIndex == -1 ? suggestion : suggestion.substring(0, outPutIndex)
    } else {
        return suggestion.substring(0, outPutIndex)
    }

}

// Function to send request to OpenAI API
api.sendToOpenAI = async (prompt, apiKey, modelType, isLastLine) => {
    if (!apiKey || !modelType) {
        alert("OpenAI API key or modelType not set.");
        return "";
    }

    const response = await fetch("https://api.openai.com/v1/completions", {
        method: "POST",
        headers: {
            "Content-Type": "application/json",
            "Authorization": `Bearer ${apiKey}`,
        },
        body: JSON.stringify({
            model: modelType,
            prompt,
            temperature: 0,
            max_tokens: 40,
        }),
        timeout: 30000
    });

    const data = await response.json();

    return data.choices && data.choices[0] ? cleanUpOpenaiOutput(data.choices[0].text, isLastLine) : ""
}


api.sendToBigcode = async (code, url, token, isLastLine) => {
    if (!url || !token) {
        alert("BigCode service URL or Huggingface Access Token not set.");
        return "";
    }

    const prompt = code.replace(/\u200B/g, '')


    const bodyData = {
        inputs: prompt,
        stream: false,
        parameters: {
            return_full_text: false,
            stop: ["<jupyter_output>"],
            max_new_tokens: 20
        },

    }

    if (!isLastLine){
        bodyData.parameters.stop.push("\n")
    }

    const response = await fetch(url, {
        method: "POST",
        headers: {
            "Content-Type": "application/json",
            "Authorization": `Bearer ${token}`,
        },
        body: JSON.stringify(bodyData)
    })

    const data = await response.json();

    return data[0] ? cleanUpBigcodeOutput(data[0].generated_text, isLastLine) : ""
}


window.api = api