diff --git a/extensions/openai/logits.py b/extensions/openai/logits.py new file mode 100644 index 0000000..9d2fe41 --- /dev/null +++ b/extensions/openai/logits.py @@ -0,0 +1,11 @@ +from extensions.openai.completions import process_parameters +from modules.logits import get_next_logits + + +def _get_next_logits(body): + # Pre-process the input payload to simulate a real generation + use_samplers = body['use_samplers'] + state = process_parameters(body) if use_samplers else {} + state['stream'] = True + + return get_next_logits(body['prompt'], state, use_samplers, "", return_dict=True)