mirror of
https://github.com/sunnypilot/sunnypilot.git
synced 2026-09-12 11:03:43 +08:00
github-actions[bot]
0e51ecbb7e
version: sunnypilot v2026.003.000 (dev)
date: 2026-09-09T15:43:49
master commit: b255314f5a
25 lines
849 B
Python
25 lines
849 B
Python
import unittest
|
|
from test.helpers import slow
|
|
from examples.mamba import Mamba, generate
|
|
from transformers import AutoTokenizer
|
|
|
|
PROMPT = 'Why is gravity '
|
|
TOKENIZER = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
|
|
|
|
@slow
|
|
class TestMamba(unittest.TestCase):
|
|
def test_mamba_130M(self):
|
|
OUT_130M = '''Why is gravity \nnot a good idea?\n\nA:'''
|
|
model = Mamba.from_pretrained('130m')
|
|
tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
|
|
self.assertEqual(OUT_130M, tinyoutput)
|
|
del model
|
|
def test_mamba_370M(self):
|
|
OUT_370M = '''Why is gravity \nso important?\nBecause it's the only'''
|
|
model = Mamba.from_pretrained('370m')
|
|
tinyoutput = generate(model, TOKENIZER, PROMPT, n_tokens_to_gen=10)
|
|
self.assertEqual(OUT_370M, tinyoutput)
|
|
del model
|
|
if __name__ == '__main__':
|
|
unittest.main()
|