Coverage for src/certus/interface.py: 100%

1"""Module for the log probability interface functions.""" 

2 

3from google.genai import types 

4 

5from . import utils 

6from .nodes import core 

7 

8 

9def from_google(result: types.LogprobsResult) -> list[core.Token]: 

10 """ 

11 Extract token nodes from a Google GenAI log-probs result. 

12 

13 Parameters 

14 ---------- 

15 result : google.genai.types.LogprobsResult 

16 Log-probs result. 

17 

18 Returns 

19 ------- 

20 list of certus.nodes.Token 

21 Token nodes. 

22 """ 

23 if result.chosen_candidates is None: 

24 return [] 

25 

26 tokens, position = [], 0 

27 for candidate in result.chosen_candidates: 

28 value = candidate.token 

29 logprob = candidate.log_probability 

30 if value is None or logprob is None: 

31 continue 

32 

33 tokens.append(core.Token(value, utils.clamp(logprob, upper=0.0), position)) 

34 position += len(value) 

35 

36 return tokens