llms-from-scratch-cn/Model_Architecture_Discussions/mamba/demo.ipynb

(1) Load model

In [1]:

%env HF_ENDPOINT=https://hf-mirror.com

env: HF_ENDPOINT=https://hf-mirror.com

In [3]:

%pip install einops

Looking in indexes: https://pypi.tuna.tsinghua.edu.cn/simple
Collecting einops
  Downloading https://pypi.tuna.tsinghua.edu.cn/packages/44/5a/f0b9ad6c0a9017e62d4735daaeb11ba3b6c009d69a26141b258cd37b5588/einops-0.8.0-py3-none-any.whl (43 kB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 43.2/43.2 kB 395.8 kB/s eta 0:00:00m eta -:--:--
[?25hInstalling collected packages: einops
Successfully installed einops-0.8.0
Note: you may need to restart the kernel to use updated packages.

In [4]:

from model import Mamba, ModelArgs
from transformers import AutoTokenizer

# One of:
#     'state-spaces/mamba-2.8b-slimpj'
#     'state-spaces/mamba-2.8b'
#     'state-spaces/mamba-1.4b'
#     'state-spaces/mamba-790m'
#     'state-spaces/mamba-370m'
#     'state-spaces/mamba-130m'
pretrained_model_name = 'state-spaces/mamba-370m'

model = Mamba.from_pretrained(pretrained_model_name)
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b')

config.json:   0%|          | 0.00/200 [00:00<?, ?B/s]

pytorch_model.bin:   0%|          | 0.00/1.49G [00:00<?, ?B/s]

/data1/ckw/micromamba/envs/kewei-ai/lib/python3.11/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()
  return self.fget.__get__(instance, owner)()

tokenizer_config.json:   0%|          | 0.00/156 [00:00<?, ?B/s]

vocab.json: 0.00B [00:00, ?B/s]

merges.txt: 0.00B [00:00, ?B/s]

tokenizer.json: 0.00B [00:00, ?B/s]

special_tokens_map.json:   0%|          | 0.00/90.0 [00:00<?, ?B/s]

(2) Generate Text

In [5]:

import torch
import torch.nn.functional as F


def generate(model,
             tokenizer,
             prompt: str,
             n_tokens_to_gen: int = 50,
             sample: bool = True,
             top_k: int = 40):
    model.eval()
    
    input_ids = tokenizer(prompt, return_tensors='pt').input_ids
    
    for token_n in range(n_tokens_to_gen):
        with torch.no_grad():
            indices_to_input = input_ids
            next_token_logits = model(indices_to_input)[:, -1]
        
        probs = F.softmax(next_token_logits, dim=-1)
        (batch, vocab_size) = probs.shape
        
        if top_k is not None:
            (values, indices) = torch.topk(probs, k=top_k)
            probs[probs < values[:, -1, None]] = 0
            probs = probs / probs.sum(axis=1, keepdims=True)
        
        if sample:
            next_indices = torch.multinomial(probs, num_samples=1)
        else:
            next_indices = torch.argmax(probs, dim=-1)[:, None]
        
        input_ids = torch.cat([input_ids, next_indices], dim=1)

    output_completions = [tokenizer.decode(output.tolist()) for output in input_ids][0]
    
    return output_completions

In [6]:

print(generate(model, tokenizer, 'Mamba is the'))

Mamba is the new bestselling game from the popular mobile developers. This addictive and immersive platformer is a fun and addictive mobile game created by Nintendo that requires players to take control of a large robot monster in an environment where you control the movement of its legs and

Cell:

[Cell type raw - unsupported, skipped]

In [7]:

print(generate(model, tokenizer, '\nDataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence.'))

DataWhalechina is an organization founded at Shanghai Jiao Tong University that helps learners learn artificial intelligence. The organization helps students gain expertise to be able to lead data mining projects in companies. The organization has developed various applications to help a learning process of data mining and artificial intelligence. Moreover, data management and the use of Big Data can also be used as

In [8]:

print(generate(model, tokenizer, 'The meaning of life is '))

The meaning of life is 
just this: It is the best you can do.

--K.J.

And finally: How to handle your emotions. 

<|endoftext|>Q:

Error creating an EntityManager instance in JavaEE 7

This is

In [11]:

print(generate(model, tokenizer, 'def reverse_string('))

def reverse_string(text, result):
        # find the position of the start of the string.
        start = text.index(text[0:-1])
        # find the position where the string begins changing.
        end = text.index

Cell:

[Cell type raw - unsupported, skipped]