Spaces:
Runtime error
Runtime error
File size: 1,765 Bytes
1cf2abd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
from ctransformers import LLM, Config
class MockLLM(LLM):
def __init__(self):
self._config = Config()
self._llm = None
@property
def config(self):
return self._config
def tokenize(self, prompt, **kwargs):
self.tokens = prompt.split(" ")
return range(len(self.tokens))
def generate(self, tokens, **kwargs):
return tokens
def detokenize(self, tokens, decode=True):
text = " " + self.tokens[tokens[0]]
if not decode:
text = text.encode()
return text
class TestLLM:
def test_stop(self):
llm = MockLLM()
prompt = "foo bar baz lorem ipsum\ndolor"
expected = [
([], " foo bar baz lorem ipsum\ndolor"),
(["dolor "], " foo bar baz lorem ipsum\ndolor"),
(["ipsum "], " foo bar baz lorem ipsum\ndolor"),
(["doloro"], " foo bar baz lorem ipsum\ndolor"),
(["ipsumo"], " foo bar baz lorem ipsum\ndolor"),
(["dolor"], " foo bar baz lorem ipsum\n"),
(["ipsum"], " foo bar baz lorem "),
(["olor"], " foo bar baz lorem ipsum\nd"),
(["olo"], " foo bar baz lorem ipsum\nd"),
(["psum"], " foo bar baz lorem i"),
(["psu"], " foo bar baz lorem i"),
(["z lor"], " foo bar ba"),
(["rem", "or"], " foo bar baz l"),
(["foo"], " "),
(["f"], " "),
([" "], ""),
(["\n"], " foo bar baz lorem ipsum"),
(["m\nd"], " foo bar baz lorem ipsu"),
]
for stop, response in expected:
assert llm(prompt, stop=stop) == response
if len(stop) == 1:
assert llm(prompt, stop=stop[0]) == response
|