| 22 | return 2 ** (-avg_logprob / math.log(2)) |
| 23 | |
| 24 | def call_codex(code_str, save_probs): |
| 25 | eos_code_str = endoftext_token + code_str |
| 26 | # engine: 'davinci-codex' is currently the best codex model |
| 27 | # max_tokens=0 means that we don't want the model to generate additional tokens |
| 28 | # logprobs=0 means that we don't want the logprobs of the alternative tokens, only the actual tokens |
| 29 | # echo=True means that we want the model to echo our prompt, in addition to our (not existing) completion |
| 30 | completion = openai.Completion.create(engine="davinci-codex", prompt=eos_code_str, |
| 31 | max_tokens=0, |
| 32 | temperature=0.0, |
| 33 | logprobs=0, |
| 34 | n=1, |
| 35 | echo=True) |
| 36 | |
| 37 | c = completion.choices[0] |
| 38 | # skipping the <|endoftext|> token |
| 39 | sum_logprobs = sum(c.logprobs.token_logprobs[1:]) |
| 40 | num_tokens = len(c.logprobs.token_logprobs[1:]) |
| 41 | if save_probs: |
| 42 | saved_probs = { |
| 43 | 'text': code_str, |
| 44 | 'tokens': c.logprobs.tokens[1:], |
| 45 | 'logprobs': c.logprobs.token_logprobs[1:], |
| 46 | 'sum_logprobs': sum_logprobs |
| 47 | } |
| 48 | else: |
| 49 | saved_probs = None |
| 50 | |
| 51 | return sum_logprobs, num_tokens, saved_probs |
| 52 | |
| 53 | if __name__ == '__main__': |
| 54 | parser = argparse.ArgumentParser() |