run format.sh

vllm-project · Jul 23, 2024 · c49532c · c49532c
1 parent 189b100
commit c49532c
Show file tree

Hide file tree

Showing 3 changed files with 195 additions and 156 deletions.
diff --git a/examples/offline_inference_chat.py b/examples/offline_inference_chat.py
@@ -14,63 +14,37 @@ def print_outputs(outputs):
 
 print("=" * 80)
 
-# In this script, we demonstrate two ways to pass input to the chat method of the LLM class:
+# In this script, we demonstrate two ways to pass input to the chat method:
 
 # Conversation with a list of dictionaries
 conversation = [
+    {"role": "system", "content": "You are a helpful assistant"},
+    {"role": "user", "content": "Hello"},
+    {"role": "assistant", "content": "Hello! How can I assist you today?"},
     {
-        'role': 'system',
-        'content': "You are a helpful assistant"
-    },
-    {
-        'role': 'user',
-        'content': "Hello"
-    },
-    {
-        'role': 'assistant',
-        'content': "Hello! How can I assist you today?"
-    },
-    {
-        'role': 'user',
-        'content': "Write an essay about the importance of higher education."
+        "role": "user",
+        "content": "Write an essay about the importance of higher education.",
     },
 ]
-outputs = llm.chat(conversation,
-                   sampling_params=sampling_params,
-                   use_tqdm=False)
+outputs = llm.chat(
+    conversation, sampling_params=sampling_params, use_tqdm=False
+)
 print_outputs(outputs)
 
-# Multiple conversations 
+# Multiple conversations
 conversations = [
     [
-        {
-            'role': 'system',
-            'content': "You are a helpful assistant"
-        },
-        {
-            'role': 'user',
-            'content': "What is dark matter?"
-        },
+        {"role": "system", "content": "You are a helpful assistant"},
+        {"role": "user", "content": "What is dark matter?"},
     ],
     [
+        {"role": "system", "content": "You are a helpful assistant"},
+        {"role": "user", "content": "How are you?"},
         {
-            'role': 'system',
-            'content': "You are a helpful assistant"
-        },
-        {
-            'role': 'user',
-            'content': "How are you?"
-        },
-        {
-            'role':
-            'assistant',
-            'content':
-            "I'm an AI, so I don't have feelings, but I'm here to help you!"
-        },
-        {
-            'role': 'user',
-            'content': "Tell me a joke."
+            "role": "assistant",
+            "content": "I'm an AI without feelings, but I'm here to help!",
         },
+        {"role": "user", "content": "Tell me a joke."},
     ],
 ]
 

diff --git a/tests/entrypoints/llm/test_generate.py b/tests/entrypoints/llm/test_generate.py
@@ -141,6 +141,7 @@ def test_multiple_sampling_params(llm: LLM):
     outputs = llm.generate(PROMPTS, sampling_params=None)
     assert len(PROMPTS) == len(outputs)
 
+
 def test_chat():
 
     llm = LLM(model=MODEL_NAME)
@@ -160,7 +161,7 @@ def test_chat():
     assert len(outputs) == 1
 
     prompt2 = "Describe Bangkok in 150 words."
-    messages = [messages] + [[
+    multiple_messages = [messages] + [[
         {
             "role": "system",
             "content": "You are a helpful assistant"
@@ -170,8 +171,8 @@ def test_chat():
             "content": prompt2
         },
     ]]
-    outputs = llm.chat(messages)
-    assert len(outputs) == len(messages)
+    outputs = llm.chat(multiple_messages)
+    assert len(outputs) == len(multiple_messages)
 
     sampling_params = [
         SamplingParams(temperature=0.01, top_p=0.95),