Support embedding input as a list (#1014)

2024-08-10 08:39:05 -07:00
parent 62757db6f0
commit 7599badeaf
3 changed files with 64 additions and 49 deletions
--- a/python/sglang/test/runners.py
+++ b/python/sglang/test/runners.py
@@ -219,11 +219,9 @@ class SRTRunner:
                output_strs=output_strs, top_input_logprobs=top_input_logprobs
            )
        else:
-            logits = []
-            for prompt in prompts:
-                response = self.runtime.encode(prompt)
-                response = json.loads(response)
-                logits.append(response["embedding"])
+            response = self.runtime.encode(prompts)
+            response = json.loads(response)
+            logits = [x["embedding"] for x in response]
            return ModelOutput(embed_logits=logits)

    def __enter__(self):