refactor for useage on runpod hub

.github/workflows/build-and-push.yml

@@ -38,7 +38,7 @@ jobs:
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - uses: docker/build-push-action@v6
         with:
-          context: ./wrapper
+          context: ./.runpod
           build-args: OLLAMA_VERSION=${{ inputs.ollama_version }}
           push: true
           tags: svenbrnn/runpod-ollama:${{ inputs.ollama_version }}${{ inputs.next_version_suffix }},svenbrnn/runpod-ollama:latest

.runpod/Dockerfile

@@ -32,7 +32,7 @@ RUN apt-get update --yes --quiet && DEBIAN_FRONTEND=noninteractive apt-get insta
 WORKDIR /work
 
 # Add my src as /work
-ADD src /work
+ADD engine.py handler.py requirements.txt start.sh test_input.json utils.py /work
 
 # Set defaut ollama models directory to /runpod-volume where runpod will mount the volume by default
 ENV OLLAMA_MODELS="/runpod-volume"

.runpod/start.sh

@@ -37,4 +37,4 @@ else
     ollama pull $MODEL_NAME
 fi
 
-python -u runpod_wrapper.py $1
+python -u handler.py $1

README.md

@@ -5,6 +5,8 @@
 Start a runpod serverless with the docker container ``svenbrnn/runpod-ollama:latest``. Set ``MODEL_NAME`` environment to a model from ollama.com to automatically download a model.
 A mounted volume will be automatically used.
 
+[![RunPod](https://api.runpod.io/badge/SvenBrnn/runpod-worker-ollama)](https://www.runpod.io/console/hub/SvenBrnn/runpod-worker-ollama)
+
 ## Environment variables
 
 | Variable Name | Description                              | Default Value       |