From 106c9ce8e5a2f69b1357140f4c4607aaf646ff41 Mon Sep 17 00:00:00 2001
From: drbh <david.richard.holtz@gmail.com>
Date: Wed, 3 Apr 2024 20:18:06 +0000
Subject: [PATCH] fix: avoid old change

---
 router/src/lib.rs | 32 +++++++++++++++++++++-----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/router/src/lib.rs b/router/src/lib.rs
index e2f21bc7..020117c0 100644
--- a/router/src/lib.rs
+++ b/router/src/lib.rs
@@ -399,23 +399,33 @@ impl From<(Token, Vec<Token>)> for ChatCompletionLogprobs {
 impl From<(Vec<Token>, Vec<Vec<Token>>)> for ChatCompletionLogprobs {
     fn from(value: (Vec<Token>, Vec<Vec<Token>>)) -> Self {
         let (tokens, top_tokens) = value;
-        Self {
-            content: tokens
-                .into_iter()
-                .zip(top_tokens)
-                .map(|(t, top_t)| ChatCompletionLogprob {
-                    token: t.text,
-                    logprob: t.logprob,
-                    top_logprobs: top_t
+
+        // Create an iterator that produces None for top_tokens once it's exhausted
+        let top_tokens_iter = top_tokens
+            .into_iter()
+            .map(Some)
+            .chain(std::iter::repeat(None));
+
+        let content = tokens
+            .into_iter()
+            .zip(top_tokens_iter)
+            .map(|(t, top_t_option)| ChatCompletionLogprob {
+                token: t.text,
+                logprob: t.logprob,
+                top_logprobs: match top_t_option {
+                    Some(top_t) => top_t
                         .into_iter()
                         .map(|t| ChatCompletionTopLogprob {
                             token: t.text,
                             logprob: t.logprob,
                         })
                         .collect(),
-                })
-                .collect(),
-        }
+                    None => vec![], // Handle the case where there are no top tokens
+                },
+            })
+            .collect();
+
+        Self { content }
     }
 }