Add session context manager for Chat pipeline

[rahul-tuli]

Adds a session context manager for chat pipelines

Tests: Added automated tests

Local Test:

from deepsparse import Pipeline

model = "/home/rahul/.cache/sparsezoo/neuralmagic/codegen_mono-350m-bigpython_bigquery_thepile-base/deployment"

pipeline = Pipeline.create(
    task="chat",
    model_path=model,
)
print("Pipeline created")
inp = {"sequences" : "def fib("}
with pipeline.session():   
    for _ in range(2):
        out = pipeline(**inp)
        print(out)
    
assert len(pipeline.engine.kv_cache_storage._memory) == 1, "Pipeline created more than one session ids"

print("Output from OG pipeline:")
print(pipeline(**inp))

assert len(pipeline.engine.kv_cache_storage._memory) == 2, "Pipeline should have 2 session ids at this level"

Output:

2023-09-14 15:00:14 deepsparse.transformers WARNING  The neuralmagic fork of transformers may not be installed. It can be installed via `pip install nm_transformers`
Using pad_token, but it is not set yet.
2023-09-14 15:00:17 deepsparse.transformers.pipelines.text_generation INFO     Compiling an auxiliary engine to process a prompt with a larger processing length. This improves performance, but may result in additional memory consumption.
2023-09-14 15:00:18 deepsparse.utils.onnx INFO     Overwriting in-place the input shapes of the transformer model at /home/rahul/.cache/sparsezoo/neuralmagic/codegen_mono-350m-bigpython_bigquery_thepile-base/deployment/model.onnx
DeepSparse, Copyright 2021-present / Neuralmagic, Inc. version: 1.6.0.20230910 COMMUNITY | (c66b57da) (release) (optimized) (system=avx512, binary=avx512)
2023-09-14 15:00:27 deepsparse.utils.onnx INFO     Overwriting in-place the input shapes of the transformer model at /home/rahul/.cache/sparsezoo/neuralmagic/codegen_mono-350m-bigpython_bigquery_thepile-base/deployment/model.onnx
Pipeline created
sequences=['n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib(n-1) + fib(n-2)\n\ndef fib2(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib2(n-1) + fib2(n-2)\n\ndef fib3(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib3(n-1) + fib3(n-2)\n\ndef fib4(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib4(n-1) + fib4(n-2)\n\ndef fib5(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib5(n-1) + fib5(n-2)\n\ndef fib6(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib6(n-1) + fib6(n-2)\n\ndef fib7(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib7(n-1) + fib7(n-2)\n\ndef fib8(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib8(n-1) + fib8(n-2)\n\ndef fib9(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib9(n-1) + fib9(n-2)\n\ndef fib10(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib10(n-1) + fib10(n-2)\n\ndef fib11(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib11(n-1) + fib11(n-2)\n\ndef fib12(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib12(n-1) + fib12(n-2)\n\ndef fib13(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib13(n-1) + fib13(n-2)\n\ndef fib14(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib14(n-1) + fib14(n-2)\n\ndef fib15(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib15(n-1) + fib15(n-2)\n\ndef fib16(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib16(n-1) + fib16(n-2)\n\ndef fib17(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib17(n-1) + fib17(n-2)\n\ndef fib18(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib18(n-1) + fib18(n-2)\n\ndef fib19(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib19(n-1) + fib19(n-2)\n\ndef fib20(n):\n    if n == 0:\n        return 0\n    elif n =='] logits=None session_ids=['09aabad9-597a-416b-a891-f5fe66e10040']
sequences=['n-1):\n        return 1\n    else:\n        return fib20(n-1) + fib20(n-2)\n\ndef fib21(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib21(n-1) + fib21(n-2)\n\ndef fib22(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib22(n-1) + fib22(n-2)\n\ndef fib23(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib23(n-1) + fib23(n-2)\n\ndef fib24(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib24(n-1) + fib24(n-2)\n\ndef fib25(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib25(n-1) + fib25(n-2)\n\ndef fib26(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib26(n-1) + fib26(n-2)\n\ndef fib27(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib27(n-1) + fib27(n-2)\n\ndef fib28(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib28(n-1) + fib28(n-2)\n\ndef fib29(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib29(n-1) + fib29(n-2)\n\ndef fib30(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib30(n-1) + fib30(n-2)\n\ndef fib31(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib31(n-1) + fib31(n-2)\n\ndef fib32(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib32(n-1) + fib32(n-2)\n\ndef fib33(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib33(n-1) + fib33(n-2)\n\ndef fib34(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib34(n-1) + fib34(n-2)\n\ndef fib35(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib35(n-1) + fib35(n-2)\n\ndef fib36(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib36(n-1) + fib36(n-2)\n\ndef fib37(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib37(n-1) + fib37(n-2)\n\ndef fib38(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib38(n-1) + fib38(n-2)\n\ndef fib39(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        \n        \n'] logits=None session_ids=['09aabad9-597a-416b-a891-f5fe66e10040']
Output from OG pipeline:
sequences=['n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib(n-1) + fib(n-2)\n\ndef fib2(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib2(n-1) + fib2(n-2)\n\ndef fib3(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib3(n-1) + fib3(n-2)\n\ndef fib4(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib4(n-1) + fib4(n-2)\n\ndef fib5(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib5(n-1) + fib5(n-2)\n\ndef fib6(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib6(n-1) + fib6(n-2)\n\ndef fib7(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib7(n-1) + fib7(n-2)\n\ndef fib8(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib8(n-1) + fib8(n-2)\n\ndef fib9(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib9(n-1) + fib9(n-2)\n\ndef fib10(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib10(n-1) + fib10(n-2)\n\ndef fib11(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib11(n-1) + fib11(n-2)\n\ndef fib12(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib12(n-1) + fib12(n-2)\n\ndef fib13(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib13(n-1) + fib13(n-2)\n\ndef fib14(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib14(n-1) + fib14(n-2)\n\ndef fib15(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib15(n-1) + fib15(n-2)\n\ndef fib16(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib16(n-1) + fib16(n-2)\n\ndef fib17(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib17(n-1) + fib17(n-2)\n\ndef fib18(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib18(n-1) + fib18(n-2)\n\ndef fib19(n):\n    if n == 0:\n        return 0\n    elif n == 1:\n        return 1\n    else:\n        return fib19(n-1) + fib19(n-2)\n\ndef fib20(n):\n    if n == 0:\n        return 0\n    elif n =='] logits=None session_ids=['3b9e7f85-39c2-44ac-b22b-ac6529d1275f']

---

• To see the specific tasks where the Asana app for GitHub is being used, see below:
  • https://app.asana.com/0/0/1205491048843456

[rahul-tuli]

These nits should be a part of a separate PR ideally, but added them here as I was already touching this file

[bfineran]

closing with #1276