diff --git a/docker-compose.yml b/docker-compose.yml
index 57e5a9a..0b56d0a 100644
--- a/docker-compose.yml
+++ b/docker-compose.yml
@@ -4,16 +4,26 @@ services:
     # 12D3KooWM5RbfAgPAiMcXFFkWEyWBcfRigMzuxNewynSJ8YG6PsZ
     container_name: worker
     build: .
-    command: >
-      ./upshot-node --role=worker --peer-db=/data/peerdb --function-db=/data/function-db 
-      --runtime-path=/app/runtime --runtime-cli=bls-runtime --workspace=/data/workspace
-       --private-key=/data/keys/worker/priv.bin --log-level=debug --port=9011 
-       --boot-nodes=/ip4/172.19.0.100/tcp/9010/p2p/12D3KooWPaHbkXmKQLrQ8xfQj39pvbfhrQSWMqJSCAE9WYaqcE7R
+    entrypoint:
+      - "/bin/bash"
+      - "-c"
+      - |
+        if [ ! -f /data/keys/priv.bin ]; then
+          echo "Generating new private keys..."
+          mkdir -p /data/keys
+          cd /data/keys
+          allora-keys
+        fi
+        # Change boot-nodes below to the key advertised by your head
+        allora-node --role=worker --peer-db=/data/peerdb --function-db=/data/function-db \
+          --runtime-path=/app/runtime --runtime-cli=bls-runtime --workspace=/data/workspace \
+          --private-key=/data/keys/priv.bin --log-level=debug --port=9011 \
+          --boot-nodes=/ip4/172.19.0.100/tcp/9010/p2p/12D3KooWEUkndnKPD3AVxxVrks6ibTUTbxHe1uMYsCSNTzGJv6N9 \
+          --topic=2
     volumes:
       - type: bind
-        source: ./keys/worker1/
-        target: /data/keys/worker
-        read_only: true
+        source: ./worker-data
+        target: /data
     env_file:
       - .env
     depends_on:
@@ -25,20 +35,28 @@ services:
         ipv4_address: 172.19.0.5
 
   head:
-    # 12D3KooWPaHbkXmKQLrQ8xfQj39pvbfhrQSWMqJSCAE9WYaqcE7R
+    # 12D3KooWEUkndnKPD3AVxxVrks6ibTUTbxHe1uMYsCSNTzGJv6N9
     container_name: head
     image: alloranetwork/allora-inference-base-head:v0.0.5
-    command: >
-      ./upshot-node --role=head --peer-db=/data/peerdb --function-db=/data/function-db
-      --runtime-path=/app/runtime --runtime-cli=bls-runtime --workspace=/data/workspace
-      --private-key=/data/keys/head/priv.bin --log-level=debug --port=9010 --rest-api=:6000
+    environment:
+      - HOME=/data
+    entrypoint:
+      - "/bin/bash"
+      - "-c"
+      - |
+        if [ ! -f /data/keys/priv.bin ]; then
+          echo "Generating new private keys..."
+          mkdir -p /data/keys
+          cd /data/keys
+          allora-keys
+        fi
+        allora-node --role=head --peer-db=/data/peerdb --function-db=/data/function-db  \
+          --runtime-path=/app/runtime --runtime-cli=bls-runtime --workspace=/data/workspace \
+          --private-key=/data/keys/priv.bin --log-level=debug --port=9010 --rest-api=:6000
     ports:
       - "6000:6000"
     volumes:
-      - type: bind
-        source: ./keys/head1/
-        target: /data/keys/head
-        read_only: true
+      - ./head-data:/data
 
     networks:
       b7s-local:
@@ -53,3 +71,6 @@ networks:
       config:
         - subnet: 172.19.0.0/24
 
+volumes:
+  worker-data:
+  head-data:
\ No newline at end of file
diff --git a/main.py b/main.py
index a8fba45..be410f0 100644
--- a/main.py
+++ b/main.py
@@ -47,14 +47,17 @@ def process(index_name):
         response = json.dumps({"value": str(wei_value)}) 
     else:
         response = json.dumps({"error": "No predictions"})
-    print(response)
+    return response
 
 
 if __name__ == "__main__":
     # Your code logic with the parsed argument goes here
+    # Parameters: topic_id, index_name
     try:
-        index_name = sys.argv[1]
-        process(index_name)
+        # Not using to discriminate by topicId for simplicity.
+        # topic_id = sys.argv[1]
+        index_name = sys.argv[2]
+        response = process(index_name)
     except Exception as e:
         response = json.dumps({"error": {str(e)}})
-        print(response)
+    print(response)