Spaces:
Runtime error
Runtime error
from ctransformers import LLM, Config | |
class MockLLM(LLM): | |
def __init__(self): | |
self._config = Config() | |
self._llm = None | |
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 | |