mamba-chat.ipynb

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### Install dependencies and load Libraries","metadata":{}},{"cell_type":"code","source":"!pip install -Uqqq pip --progress-bar off\n!pip install -qqq torch==2.1.2 --progress-bar off\n!pip install -qqq causal-conv1d==1.1.1 --progress-bar off\n!pip install -qqq mamba-ssm==1.1.1 --progress-bar off","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:28:15.167725Z","iopub.execute_input":"2024-03-02T19:28:15.168098Z","iopub.status.idle":"2024-03-02T19:29:05.126793Z","shell.execute_reply.started":"2024-03-02T19:28:15.168067Z","shell.execute_reply":"2024-03-02T19:29:05.125654Z"},"trusted":true},"execution_count":2,"outputs":[]},{"cell_type":"code","source":"from inspect import cleandoc\nimport torch\nfrom mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel\nfrom transformers import AutoTokenizer\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nDEVICE","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:29:13.320185Z","iopub.execute_input":"2024-03-02T19:29:13.320878Z","iopub.status.idle":"2024-03-02T19:29:13.372089Z","shell.execute_reply.started":"2024-03-02T19:29:13.320845Z","shell.execute_reply":"2024-03-02T19:29:13.371121Z"},"trusted":true},"execution_count":3,"outputs":[{"execution_count":3,"output_type":"execute_result","data":{"text/plain":"'cuda'"},"metadata":{}}]},{"cell_type":"markdown","source":"### Load model + Tokenizer","metadata":{}},{"cell_type":"code","source":"model_name = \"havenhq/mamba-chat\"\nANSWER_START = \"<|assistant|>\\n\"\nANSWER_END = \"<|endoftext|>\"\n\ntokenizer = AutoTokenizer.from_pretrained(model_name)\ntokenizer.eos_token = \"<|endoftext|>\"\ntokenizer.pad_token = tokenizer.eos_token\ntokenizer.chat_template = AutoTokenizer.from_pretrained(\n    \"HuggingFaceH4/zephyr-7b-beta\"\n).chat_template\n\nmodel = MambaLMHeadModel.from_pretrained(\n    model_name,\n    device = DEVICE,\n    dtype = torch.float16\n)","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:35:18.585270Z","iopub.execute_input":"2024-03-02T19:35:18.586014Z","iopub.status.idle":"2024-03-02T19:35:53.918983Z","shell.execute_reply.started":"2024-03-02T19:35:18.585981Z","shell.execute_reply":"2024-03-02T19:35:53.918148Z"},"trusted":true},"execution_count":4,"outputs":[{"output_type":"display_data","data":{"text/plain":"tokenizer_config.json:   0%|          | 0.00/4.79k [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"9f31365c7ea44e33902308c66e54fd1d"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"tokenizer.json:   0%|          | 0.00/2.11M [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"2b9c6c78ad8f468b8451f7a6607018bf"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"special_tokens_map.json:   0%|          | 0.00/131 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"fa0947246c354b0dbccb2c631ef0418e"}},"metadata":{}},{"name":"stderr","text":"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"tokenizer_config.json:   0%|          | 0.00/1.43k [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"4fb7e7fd414d4a18ad5152a4046e659a"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"tokenizer.model:   0%|          | 0.00/493k [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"606ff73249234672879fc426a0f8ff37"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"tokenizer.json:   0%|          | 0.00/1.80M [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"944aaaca151c404dbc13e2e2648752b0"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"added_tokens.json:   0%|          | 0.00/42.0 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"1cfb386919af49e29ffdeca7b51ac3ae"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"special_tokens_map.json:   0%|          | 0.00/168 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"05395dacfecc40c79703f96e7aa16719"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"config.json:   0%|          | 0.00/201 [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"c8e25549743d42d28258d9dda0baf970"}},"metadata":{}},{"output_type":"display_data","data":{"text/plain":"pytorch_model.bin:   0%|          | 0.00/5.55G [00:00<?, ?B/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"66894c85fa794e9883c72f286a9c2e2b"}},"metadata":{}}]},{"cell_type":"code","source":"ANSWER_START = \"<|assistant|>\\n\"\nANSWER_END = \"<|endoftext|>\"","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:48:59.979411Z","iopub.execute_input":"2024-03-02T19:48:59.980143Z","iopub.status.idle":"2024-03-02T19:48:59.984532Z","shell.execute_reply.started":"2024-03-02T19:48:59.980110Z","shell.execute_reply":"2024-03-02T19:48:59.983457Z"},"trusted":true},"execution_count":13,"outputs":[]},{"cell_type":"markdown","source":"### Prompt + Generating Output","metadata":{}},{"cell_type":"code","source":"prompt = 'What is the capital of Egpyt?'\n\nmessages = [\n    {\n        \"role\": \"user\",\n        \"content\": prompt\n    }\n]\nmessages","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:36:47.967870Z","iopub.execute_input":"2024-03-02T19:36:47.968658Z","iopub.status.idle":"2024-03-02T19:36:47.975187Z","shell.execute_reply.started":"2024-03-02T19:36:47.968628Z","shell.execute_reply":"2024-03-02T19:36:47.974186Z"},"trusted":true},"execution_count":5,"outputs":[{"execution_count":5,"output_type":"execute_result","data":{"text/plain":"[{'role': 'user', 'content': 'What is the capital of Egpyt?'}]"},"metadata":{}}]},{"cell_type":"code","source":"input_ids = tokenizer.apply_chat_template(\n    messages, \n    return_tensors = \"pt\",\n    add_generation_prompt = True\n).to(DEVICE)\n\ninput_ids","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:38:42.090219Z","iopub.execute_input":"2024-03-02T19:38:42.090587Z","iopub.status.idle":"2024-03-02T19:38:42.139828Z","shell.execute_reply.started":"2024-03-02T19:38:42.090557Z","shell.execute_reply":"2024-03-02T19:38:42.138973Z"},"trusted":true},"execution_count":6,"outputs":[{"execution_count":6,"output_type":"execute_result","data":{"text/plain":"tensor([[   29,    93,  4537, 49651,   187,  1276,   310,   253,  5347,   273,\n           444, 17788,  1767,    32,     0,   187,    29,    93,   515,  5567,\n         49651,   187]], device='cuda:0')"},"metadata":{}}]},{"cell_type":"code","source":"outputs = model.generate(\n    input_ids = input_ids,\n    max_length = 1024,\n    temperature = 0.7,\n    top_p = 0.7,\n    eos_token_id = tokenizer.eos_token_id\n)\n\nresponse = tokenizer.decode(outputs[0])\nprint(response)","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:40:46.013416Z","iopub.execute_input":"2024-03-02T19:40:46.014288Z","iopub.status.idle":"2024-03-02T19:40:54.143650Z","shell.execute_reply.started":"2024-03-02T19:40:46.014251Z","shell.execute_reply":"2024-03-02T19:40:54.142679Z"},"trusted":true},"execution_count":7,"outputs":[{"name":"stdout","text":"<|user|>\nWhat is the capital of Egpyt?<|endoftext|>\n<|assistant|>\nThe capital of Egypt is Cairo.<|endoftext|>\n","output_type":"stream"}]},{"cell_type":"markdown","source":"### Testing","metadata":{}},{"cell_type":"code","source":"def predict(prompt: str, system_prompt: str = \"\") -> str:\n    messages = []\n    \n    if system_prompt:\n        messages.append(\n            {\n                \"role\": \"system\",\n                \"content\": system_prompt\n            }\n        )\n        \n    messages.append(\n            {\n                \"role\": \"user\",\n                \"content\": prompt\n            }\n        )\n    \n    input_ids = tokenizer.apply_chat_template(\n        messages,\n        return_tensors = \"pt\",\n        add_generation_prompt = True\n    ).to(DEVICE)\n    \n    outputs = model.generate(\n        input_ids = input_ids,\n        max_length = 1024,\n        temperature = 0.9,\n        top_p = 0.7,\n        eos_token_id = tokenizer.eos_token_id\n    )\n    \n    response = tokenizer.decode(outputs[0])\n    \n    return extract_response(response)","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:52:46.889709Z","iopub.execute_input":"2024-03-02T19:52:46.890579Z","iopub.status.idle":"2024-03-02T19:52:46.897420Z","shell.execute_reply.started":"2024-03-02T19:52:46.890547Z","shell.execute_reply":"2024-03-02T19:52:46.896513Z"},"trusted":true},"execution_count":20,"outputs":[]},{"cell_type":"code","source":"def extract_response(output: str) -> str:\n    response_start = output.find(ANSWER_START) + len(ANSWER_START)\n    return output[response_start : output.find(ANSWER_END, response_start)]","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:49:05.091215Z","iopub.execute_input":"2024-03-02T19:49:05.092072Z","iopub.status.idle":"2024-03-02T19:49:05.096729Z","shell.execute_reply.started":"2024-03-02T19:49:05.092042Z","shell.execute_reply":"2024-03-02T19:49:05.095740Z"},"trusted":true},"execution_count":14,"outputs":[]},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nWhat is the capital of Egypt and what is famous for?\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:53:54.941517Z","iopub.execute_input":"2024-03-02T19:53:54.942232Z","iopub.status.idle":"2024-03-02T19:53:56.637268Z","shell.execute_reply.started":"2024-03-02T19:53:54.942198Z","shell.execute_reply":"2024-03-02T19:53:56.636273Z"},"trusted":true},"execution_count":23,"outputs":[{"name":"stdout","text":"The capital of Egypt is Cairo. It is famous for its ancient monuments, museums, and historical sites.\n\nCPU times: user 1.69 s, sys: 1.34 ms, total: 1.69 s\nWall time: 1.69 s\n","output_type":"stream"}]},{"cell_type":"markdown","source":"##### Coding","metadata":{}},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nWrite a python function that calculates the squareroot of a multiplication of two numbers.\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:53:21.976062Z","iopub.execute_input":"2024-03-02T19:53:21.976985Z","iopub.status.idle":"2024-03-02T19:53:39.359373Z","shell.execute_reply.started":"2024-03-02T19:53:21.976952Z","shell.execute_reply":"2024-03-02T19:53:39.358312Z"},"trusted":true},"execution_count":22,"outputs":[{"name":"stdout","text":"Here's a Python function that calculates the squareroot of a multiplication of two numbers:\n\n```python\ndef squareroot(num1, num2):\n    \"\"\"\n    Calculates the squareroot of a multiplication of two numbers.\n    \"\"\"\n    \n    # Check if both numbers are positive\n    if num1 < 0 or num2 < 0:\n        raise ValueError(\"Both numbers must be positive\")\n    \n    # Calculate the square root of the product\n    sqrt_num = sqrt(num1 * num2)\n    \n    # Return the square root of the product\n    return sqrt_num\n```\n\nHere's how you can use the function:\n\n```python\nnum1 = 5\nnum2 = 3\n\nprint(squareroot(num1, num2)) # Output: 2.0\n```\n\nIn this example, the function is called with the multiplication of `num1` and `num2` as arguments. The function then calculates the square root of the product and returns it. The output of the function is `2.0`.\n\nCPU times: user 17.4 s, sys: 858 µs, total: 17.4 s\nWall time: 17.4 s\n","output_type":"stream"}]},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nWrite a function in python that checks for palindrom number.\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:54:53.212595Z","iopub.execute_input":"2024-03-02T19:54:53.213338Z","iopub.status.idle":"2024-03-02T19:55:18.844419Z","shell.execute_reply.started":"2024-03-02T19:54:53.213305Z","shell.execute_reply":"2024-03-02T19:55:18.843458Z"},"trusted":true},"execution_count":24,"outputs":[{"name":"stdout","text":"Here's a function in Python that checks for palindrom number:\n\n```python\ndef is_palindrome(num):\n    if num == num[::-1]:\n        return True\n    else:\n        return False\n```\n\nThis function takes a number `num` as an argument and checks if it is a palindrome or not. A palindrome number is a number that is equal to its reverse. For example, `1234` is a palindrome number because it is equal to its reverse `4321`.\n\nTo check if a number is a palindrome or not, we can use the `is_palindrome` function. If the function returns `True`, then the number is a palindrome. Otherwise, it is not a palindrome.\n\nTo test the function, we can use the following code:\n\n```python\nnum = 1234\nprint(is_palindrome(num)) # True\nprint(is_palindrome('1234')) # True\nprint(is_palindrome('1234')) # False\n```\n\nIn the first line, we call the `is_palindrome` function with the number `1234`. The function returns `True` because the number is a palindrome. In the second line, we call the `is_palindrome` function with the string `1234`. The function returns `True` because the string is a palindrome. In the third line, we call the `is_palindrome` function with the number `1234`. The function returns `False` because the number is not a palindrome.\n\nCPU times: user 25.6 s, sys: 8.84 ms, total: 25.6 s\nWall time: 25.6 s\n","output_type":"stream"}]},{"cell_type":"markdown","source":"##### Math","metadata":{}},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nCalculate the answer:\n3 + 8 - 2 = ?\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:52:56.088148Z","iopub.execute_input":"2024-03-02T19:52:56.088511Z","iopub.status.idle":"2024-03-02T19:52:56.591370Z","shell.execute_reply.started":"2024-03-02T19:52:56.088482Z","shell.execute_reply":"2024-03-02T19:52:56.590478Z"},"trusted":true},"execution_count":21,"outputs":[{"name":"stdout","text":"The answer is 11.\n\nCPU times: user 495 ms, sys: 4.55 ms, total: 499 ms\nWall time: 498 ms\n","output_type":"stream"}]},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nCalculate the answer:\n3 + 15 * 3 - 4 = ?\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:57:37.863081Z","iopub.execute_input":"2024-03-02T19:57:37.863459Z","iopub.status.idle":"2024-03-02T19:57:39.448199Z","shell.execute_reply.started":"2024-03-02T19:57:37.863428Z","shell.execute_reply":"2024-03-02T19:57:39.447313Z"},"trusted":true},"execution_count":27,"outputs":[{"name":"stdout","text":"The answer is:\n3 + 15 * 3 - 4 = 15 + 45 - 4 = 34\n\nCPU times: user 1.58 s, sys: 1.44 ms, total: 1.58 s\nWall time: 1.58 s\n","output_type":"stream"}]},{"cell_type":"markdown","source":"##### Reasoning","metadata":{}},{"cell_type":"code","source":"%%time\nprompt = cleandoc(\n    \"\"\"\nNaira is faster than Aml and Aml is faster than Rehab.\nIs Rehab is faster than Naira?\n\"\"\"\n)\nprint(predict(prompt))\nprint()","metadata":{"execution":{"iopub.status.busy":"2024-03-02T19:56:35.600945Z","iopub.execute_input":"2024-03-02T19:56:35.601981Z","iopub.status.idle":"2024-03-02T19:56:36.656328Z","shell.execute_reply.started":"2024-03-02T19:56:35.601947Z","shell.execute_reply":"2024-03-02T19:56:36.655395Z"},"trusted":true},"execution_count":26,"outputs":[{"name":"stdout","text":"No, Rehab is not faster than Naira.\n\nCPU times: user 1.05 s, sys: 716 µs, total: 1.05 s\nWall time: 1.05 s\n","output_type":"stream"}]}]}