runpod serverless ollama - WIP

1 files changed, 32 insertions, 0 deletions

diff --git a/runpod/main.py b/runpod/main.py
new file mode 100755
index 0000000..ae12ba4
--- /dev/null
+++ b/runpod/main.py
@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+"""ollama-runpod endpoint"""
+
+
+import runpod
+import requests
+
+
+def ollama(job):
+    """the endpoint handler"""
+    headers = {"Content-Type": "application/json"}
+    data = {
+        "model": job["model"],
+        "system": job["prompt_system"],
+        "prompt": job["prompt_user"] + job["content"],
+        "stream": False,
+        "format": "json",
+        "options": {
+            "temperature": job["temp"],
+        },
+    }
+
+    response = requests.post(
+        "http://127.0.0.1:11434/api/generate", data=data, headers=headers, timeout=10
+    )
+
+    return response
+
+
+def main() -> None:
+    """entrypoint"""
+    runpod.serverless({"handler": ollama})