add insert support to generate endpoint (#215)

* add suffix * update fill-in-the-middle example * keep example * lint * variables
ollama · Jul 18, 2024 · 33c4b61 · 33c4b61
1 parent b0ea6d9
commit 33c4b61
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/examples/fill-in-middle/main.py b/examples/fill-in-middle/main.py
@@ -1,16 +1,16 @@
 from ollama import generate
 
-prefix = '''def remove_non_ascii(s: str) -> str:
+prompt = '''def remove_non_ascii(s: str) -> str:
     """ '''
 
 suffix = """
     return result
 """
 
-
 response = generate(
   model='codellama:7b-code',
-  prompt=f'<PRE> {prefix} <SUF>{suffix} <MID>',
+  prompt=prompt,
+  suffix=suffix,
   options={
     'num_predict': 128,
     'temperature': 0,

diff --git a/ollama/_client.py b/ollama/_client.py
@@ -102,6 +102,7 @@ def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -118,6 +119,7 @@ def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -133,6 +135,7 @@ def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -162,6 +165,7 @@ def generate(
       json={
         'model': model,
         'prompt': prompt,
+        'suffix': suffix,
         'system': system,
         'template': template,
         'context': context or [],
@@ -518,6 +522,7 @@ async def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -534,6 +539,7 @@ async def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -549,6 +555,7 @@ async def generate(
     self,
     model: str = '',
     prompt: str = '',
+    suffix: str = '',
     system: str = '',
     template: str = '',
     context: Optional[Sequence[int]] = None,
@@ -577,6 +584,7 @@ async def generate(
       json={
         'model': model,
         'prompt': prompt,
+        'suffix': suffix,
         'system': system,
         'template': template,
         'context': context or [],

diff --git a/tests/test_client.py b/tests/test_client.py
@@ -137,6 +137,7 @@ def test_client_generate(httpserver: HTTPServer):
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],
@@ -182,6 +183,7 @@ def generate():
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],
@@ -210,6 +212,7 @@ def test_client_generate_images(httpserver: HTTPServer):
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],
@@ -619,6 +622,7 @@ async def test_async_client_generate(httpserver: HTTPServer):
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],
@@ -659,6 +663,7 @@ def generate():
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],
@@ -688,6 +693,7 @@ async def test_async_client_generate_images(httpserver: HTTPServer):
     json={
       'model': 'dummy',
       'prompt': 'Why is the sky blue?',
+      'suffix': '',
       'system': '',
       'template': '',
       'context': [],