Views
No views yet
1import torch
2from transformers import AutoTokenizer
3from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
4
5device = "cuda"
6tokenizer = AutoTokenizer.from_pretrained("Schmadge/mamba-slim-orca")
7tokenizer.eos_token = tokenizer.pad_token = "<|endoftext|>"
8tokenizer.chat_template = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta").chat_template
9model = MambaLMHeadModel.from_pretrained("Schmadge/mamba-slim-orca", device=device, dtype=torch.float16)
10
11def generate_response(system_prompt, user_prompt):
12 # Preparing the prompt
13 prompt = [
14 {"role": "system", "content": system_prompt},
15 {"role": "user", "content": user_prompt}
16 ]
17 input_ids = tokenizer.apply_chat_template(prompt, return_tensors="pt", add_generation_prompt=True).to(device)
18
19 # Generating the response
20 out = model.generate(input_ids=input_ids, max_length=2000, temperature=0.3, top_p=0.7, eos_token_id=tokenizer.eos_token_id)
21 decoded = tokenizer.batch_decode(out)
22
23 return decoded[0].split("<|assistant|>\n")[-1].replace('<|endoftext|>','')
24
25system_prompt = "You are an AI assistant. Provide a detailed answer so user don't need to search outside to understand the answer."
26user_prompt = "In a room I have only 3 sisters. Anna is reading a book. Alice is playing a match of chess.What the third sister, Amanda is doing ?"
27response = generate_response(system_prompt, user_prompt)
28print(response)
29#Based on the information provided, we can infer that Amanda is playing a match of chess with Alice. Since Anna is reading a book, it is reasonable to assume that Amanda is playing a game of chess with Alice, as this is a common activity for siblings to engage in together.1@misc{haven2023mambachat,
2 title = {Mamba-Chat},
3 author = {Justus Mattern and Konstantin Hohr},
4 year = {2023},
5 howpublished = {GitHub},
6 url = {https://github.com/havenhq/mamba-chat}
7}1@article{mamba,
2 title={Mamba: Linear-Time Sequence Modeling with Selective State Spaces},
3 author={Gu, Albert and Dao, Tri},
4 journal={arXiv preprint arXiv:2312.00752},
5 year={2023}
6}1@misc{SlimOrca,
2 title = {SlimOrca: An Open Dataset of GPT-4 Augmented FLAN Reasoning Traces, with Verification},
3 author = {Wing Lian and others},
4 year = {2023},
5 publisher = {HuggingFace},
6 url = {https://huggingface.co/Open-Orca/SlimOrca}
7}