Skip to content

Commit 33c4b61

Browse files
authored
add insert support to generate endpoint (#215)
* add suffix * update fill-in-the-middle example * keep example * lint * variables
1 parent b0ea6d9 commit 33c4b61

File tree

3 files changed

+17
-3
lines changed

3 files changed

+17
-3
lines changed

examples/fill-in-middle/main.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,16 @@
11
from ollama import generate
22

3-
prefix = '''def remove_non_ascii(s: str) -> str:
3+
prompt = '''def remove_non_ascii(s: str) -> str:
44
""" '''
55

66
suffix = """
77
return result
88
"""
99

10-
1110
response = generate(
1211
model='codellama:7b-code',
13-
prompt=f'<PRE> {prefix} <SUF>{suffix} <MID>',
12+
prompt=prompt,
13+
suffix=suffix,
1414
options={
1515
'num_predict': 128,
1616
'temperature': 0,

ollama/_client.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,7 @@ def generate(
102102
self,
103103
model: str = '',
104104
prompt: str = '',
105+
suffix: str = '',
105106
system: str = '',
106107
template: str = '',
107108
context: Optional[Sequence[int]] = None,
@@ -118,6 +119,7 @@ def generate(
118119
self,
119120
model: str = '',
120121
prompt: str = '',
122+
suffix: str = '',
121123
system: str = '',
122124
template: str = '',
123125
context: Optional[Sequence[int]] = None,
@@ -133,6 +135,7 @@ def generate(
133135
self,
134136
model: str = '',
135137
prompt: str = '',
138+
suffix: str = '',
136139
system: str = '',
137140
template: str = '',
138141
context: Optional[Sequence[int]] = None,
@@ -162,6 +165,7 @@ def generate(
162165
json={
163166
'model': model,
164167
'prompt': prompt,
168+
'suffix': suffix,
165169
'system': system,
166170
'template': template,
167171
'context': context or [],
@@ -518,6 +522,7 @@ async def generate(
518522
self,
519523
model: str = '',
520524
prompt: str = '',
525+
suffix: str = '',
521526
system: str = '',
522527
template: str = '',
523528
context: Optional[Sequence[int]] = None,
@@ -534,6 +539,7 @@ async def generate(
534539
self,
535540
model: str = '',
536541
prompt: str = '',
542+
suffix: str = '',
537543
system: str = '',
538544
template: str = '',
539545
context: Optional[Sequence[int]] = None,
@@ -549,6 +555,7 @@ async def generate(
549555
self,
550556
model: str = '',
551557
prompt: str = '',
558+
suffix: str = '',
552559
system: str = '',
553560
template: str = '',
554561
context: Optional[Sequence[int]] = None,
@@ -577,6 +584,7 @@ async def generate(
577584
json={
578585
'model': model,
579586
'prompt': prompt,
587+
'suffix': suffix,
580588
'system': system,
581589
'template': template,
582590
'context': context or [],

tests/test_client.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ def test_client_generate(httpserver: HTTPServer):
137137
json={
138138
'model': 'dummy',
139139
'prompt': 'Why is the sky blue?',
140+
'suffix': '',
140141
'system': '',
141142
'template': '',
142143
'context': [],
@@ -182,6 +183,7 @@ def generate():
182183
json={
183184
'model': 'dummy',
184185
'prompt': 'Why is the sky blue?',
186+
'suffix': '',
185187
'system': '',
186188
'template': '',
187189
'context': [],
@@ -210,6 +212,7 @@ def test_client_generate_images(httpserver: HTTPServer):
210212
json={
211213
'model': 'dummy',
212214
'prompt': 'Why is the sky blue?',
215+
'suffix': '',
213216
'system': '',
214217
'template': '',
215218
'context': [],
@@ -619,6 +622,7 @@ async def test_async_client_generate(httpserver: HTTPServer):
619622
json={
620623
'model': 'dummy',
621624
'prompt': 'Why is the sky blue?',
625+
'suffix': '',
622626
'system': '',
623627
'template': '',
624628
'context': [],
@@ -659,6 +663,7 @@ def generate():
659663
json={
660664
'model': 'dummy',
661665
'prompt': 'Why is the sky blue?',
666+
'suffix': '',
662667
'system': '',
663668
'template': '',
664669
'context': [],
@@ -688,6 +693,7 @@ async def test_async_client_generate_images(httpserver: HTTPServer):
688693
json={
689694
'model': 'dummy',
690695
'prompt': 'Why is the sky blue?',
696+
'suffix': '',
691697
'system': '',
692698
'template': '',
693699
'context': [],

0 commit comments

Comments
 (0)