kanshan/ollama-python
1
0
Fork 0
mirror of https://github.com/ollama/ollama-python.git synced 2026-05-01 19:58:17 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix parse modelfile

Browse Source 
- do not add newlines while parsing
- do not add leading whitespace
This commit is contained in:
Michael Yang 2024-01-29 12:50:45 -08:00
parent f618a2f448
commit 8e5d431d0d
2 changed files with 117 additions and 14 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
ollama/_client.py
View File

@ -259,13 +259,16 @@ class Client(BaseClient):
out = io.StringIO() out = io.StringIO()
for line in io.StringIO(modelfile): for line in io.StringIO(modelfile):
command, _, args = line.partition(' ') command, _, args = line.partition(' ')
if command.upper() in ['FROM', 'ADAPTER']: if command.upper() not in ['FROM', 'ADAPTER']:
path = Path(args.strip()).expanduser() print(line, end='', file=out)
path = path if path.is_absolute() else base / path continue
if path.exists():
args = f'@{self._create_blob(path)}' path = Path(args.strip()).expanduser()
path = path if path.is_absolute() else base / path
if path.exists():
args = f'@{self._create_blob(path)}\n'
print(command, args, end='', file=out)
print(command, args, file=out)
return out.getvalue() return out.getvalue()
def _create_blob(self, path: Union[str, Path]) -> str: def _create_blob(self, path: Union[str, Path]) -> str:
@ -527,13 +530,16 @@ class AsyncClient(BaseClient):
out = io.StringIO() out = io.StringIO()
for line in io.StringIO(modelfile): for line in io.StringIO(modelfile):
command, _, args = line.partition(' ') command, _, args = line.partition(' ')
if command.upper() in ['FROM', 'ADAPTER']: if command.upper() not in ['FROM', 'ADAPTER']:
path = Path(args).expanduser() print(line, end='', file=out)
path = path if path.is_absolute() else base / path continue
if path.exists():
args = f'@{await self._create_blob(path)}' path = Path(args.strip()).expanduser()
path = path if path.is_absolute() else base / path
if path.exists():
args = f'@{await self._create_blob(path)}\n'
print(command, args, end='', file=out)
print(command, args, file=out)
return out.getvalue() return out.getvalue()
async def _create_blob(self, path: Union[str, Path]) -> str: async def _create_blob(self, path: Union[str, Path]) -> str:

101
tests/test_client.py
View File

@ -416,13 +416,61 @@ def test_client_create_modelfile(httpserver: HTTPServer):
assert isinstance(response, dict) assert isinstance(response, dict)
def test_client_create_modelfile_roundtrip(httpserver: HTTPServer):
httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='HEAD').respond_with_response(Response(status=200))
httpserver.expect_ordered_request(
'/api/create',
method='POST',
json={
'name': 'dummy',
'modelfile': '''FROM @sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
TEMPLATE """[INST] <<SYS>>{{.System}}<</SYS>>
{{.Prompt}} [/INST]"""
SYSTEM """
Use
multiline
strings.
"""
PARAMETER stop [INST]
PARAMETER stop [/INST]
PARAMETER stop <<SYS>>
PARAMETER stop <</SYS>>''',
'stream': False,
},
).respond_with_json({})
client = Client(httpserver.url_for('/'))
with tempfile.NamedTemporaryFile() as blob:
response = client.create(
'dummy',
modelfile='\n'.join(
[
f'FROM {blob.name}',
'TEMPLATE """[INST] <<SYS>>{{.System}}<</SYS>>',
'{{.Prompt}} [/INST]"""',
'SYSTEM """',
'Use',
'multiline',
'strings.',
'"""',
'PARAMETER stop [INST]',
'PARAMETER stop [/INST]',
'PARAMETER stop <<SYS>>',
'PARAMETER stop <</SYS>>',
]
),
)
assert isinstance(response, dict)
def test_client_create_from_library(httpserver: HTTPServer): def test_client_create_from_library(httpserver: HTTPServer):
httpserver.expect_ordered_request( httpserver.expect_ordered_request(
'/api/create', '/api/create',
method='POST', method='POST',
json={ json={
'name': 'dummy', 'name': 'dummy',
'modelfile': 'FROM llama2\n', 'modelfile': 'FROM llama2',
'stream': False, 'stream': False,
}, },
).respond_with_json({}) ).respond_with_json({})
@ -820,6 +868,55 @@ async def test_async_client_create_modelfile(httpserver: HTTPServer):
assert isinstance(response, dict) assert isinstance(response, dict)
@pytest.mark.asyncio
async def test_async_client_create_modelfile_roundtrip(httpserver: HTTPServer):
httpserver.expect_ordered_request(PrefixPattern('/api/blobs/'), method='HEAD').respond_with_response(Response(status=200))
httpserver.expect_ordered_request(
'/api/create',
method='POST',
json={
'name': 'dummy',
'modelfile': '''FROM @sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855
TEMPLATE """[INST] <<SYS>>{{.System}}<</SYS>>
{{.Prompt}} [/INST]"""
SYSTEM """
Use
multiline
strings.
"""
PARAMETER stop [INST]
PARAMETER stop [/INST]
PARAMETER stop <<SYS>>
PARAMETER stop <</SYS>>''',
'stream': False,
},
).respond_with_json({})
client = AsyncClient(httpserver.url_for('/'))
with tempfile.NamedTemporaryFile() as blob:
response = await client.create(
'dummy',
modelfile='\n'.join(
[
f'FROM {blob.name}',
'TEMPLATE """[INST] <<SYS>>{{.System}}<</SYS>>',
'{{.Prompt}} [/INST]"""',
'SYSTEM """',
'Use',
'multiline',
'strings.',
'"""',
'PARAMETER stop [INST]',
'PARAMETER stop [/INST]',
'PARAMETER stop <<SYS>>',
'PARAMETER stop <</SYS>>',
]
),
)
assert isinstance(response, dict)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_client_create_from_library(httpserver: HTTPServer): async def test_async_client_create_from_library(httpserver: HTTPServer):
httpserver.expect_ordered_request( httpserver.expect_ordered_request(
@ -827,7 +924,7 @@ async def test_async_client_create_from_library(httpserver: HTTPServer):
method='POST', method='POST',
json={ json={
'name': 'dummy', 'name': 'dummy',
'modelfile': 'FROM llama2\n', 'modelfile': 'FROM llama2',
'stream': False, 'stream': False,
}, },
).respond_with_json({}) ).respond_with_json({})