From f0d66cf817ed4d1698556af08369b99a2180ffe5 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sun, 19 Nov 2023 10:12:13 -0800
Subject: [PATCH] Add missing file

---
 extensions/openai/logits.py | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 extensions/openai/logits.py

diff --git a/extensions/openai/logits.py b/extensions/openai/logits.py
new file mode 100644
index 00000000..9d2fe41c
--- /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)