lint fixes

Signed-off-by: Sachin Varghese <sachin.mathew31@gmail.com>

Author: SachinVarghese

examples/vllm/vllm_server.ipynb

@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -91,14 +91,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Using configuration from: ../../config.yml\n",
+      "Using configuration from: config.yml\n",
       "Run started\n",
       "Run completed\n",
       "\n",
@@ -150,7 +150,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": ".venv",
+   "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
@@ -164,7 +164,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.8"
+   "version": "3.11.2"
   }
  },
  "nbformat": 4,

inference_perf/datagen/hf_sharegpt_datagen.py

@@ -27,22 +27,31 @@ def __init__(self) -> None:
                 split="train",
             )
         )
+        self.max_num_turns = 2
+        self.data_key = "conversations"
+        self.role_key = "from"
+        self.content_key = "value"
         # initialize data collection
         next(self.sharegpt_dataset)
 
     def get_data(self) -> Generator[InferenceData, None, None]:
         if self.sharegpt_dataset is not None:
             while True:
                 data = next(self.sharegpt_dataset)
-                if data is None or len(data["conversations"]) > 2 or len(data["conversations"]) == 0:
+                if (
+                    data is None
+                    or data[self.data_key] is None
+                    or len(data[self.data_key]) > self.max_num_turns
+                    or len(data[self.data_key]) == 0
+                ):
                     continue
                 else:
                     yield InferenceData(
                         type=APIType.Chat,
                         chat=ChatCompletionData(
                             messages=[
-                                ChatMessage(role=conversation["from"], content=conversation["value"])
-                                for conversation in data["conversations"]
+                                ChatMessage(role=conversation[self.role_key], content=conversation[self.content_key])
+                                for conversation in data[self.data_key]
                             ]
                         ),
                     )