Merge pull request #584 from UKP-SQuARE/replicate

Add support for Replicate.ai models
UKP-SQuARE · May 22, 2024 · 4abea25 · 4abea25
2 parents 5fd2d48 + a541740
commit 4abea25
Show file tree

Hide file tree

Showing 13 changed files with 1,110 additions and 107 deletions.
diff --git a/.github/workflows/build-and-deploy.yml b/.github/workflows/build-and-deploy.yml
@@ -49,6 +49,7 @@ jobs:
           "skill-kgqa-kqapro",
           "evaluator",
           "sensitivity",
+          "replicate",
         ]
         include:
           - build-args: ""
@@ -118,6 +119,8 @@ jobs:
             context: ./evaluator
           - service: sensitivity
             context: ./sensitivity
+          - service: replicate
+            context: ./replicate
     steps:
       - name: Checkout
         uses: actions/checkout@v3

diff --git a/docker-compose.ytt.yaml b/docker-compose.ytt.yaml
@@ -458,6 +458,20 @@ services:
       - "traefik.http.middlewares.sensitivity-stripprefix.stripPrefixRegex.regex=/api/[a-zA-Z0-9_-]+"
       - "traefik.http.middlewares.sensitivity-addprefix.addPrefix.prefix=/api"
 
+  replicate:
+    image: #@ "ukpsquare/replicate:" + data.values.tag
+    build: 
+      context: ./replicate
+    labels:
+      - "traefik.enable=true"
+      - "traefik.http.routers.replicate.rule=PathPrefix(`/api/replicate`)"
+      - "traefik.http.routers.replicate.entrypoints=websecure"
+      - "traefik.http.routers.replicate.tls=true"
+      - "traefik.http.routers.replicate.tls.certresolver=le"
+      - "traefik.http.routers.replicate.middlewares=replicate-stripprefix,replicate-addprefix"
+      - "traefik.http.middlewares.replicate-stripprefix.stripPrefixRegex.regex=/api/[a-zA-Z0-9_-]+"
+      - "traefik.http.middlewares.replicate-addprefix.addPrefix.prefix=/api"
+
   dpr_worker:
     image: #@ "ukpsquare/model-inference-transformer:" + data.values.tag
     build:

diff --git a/frontend/src/api/index.js b/frontend/src/api/index.js
@@ -12,7 +12,7 @@ const SKILL_MANAGER_URL = `${process.env.VUE_APP_SKILL_MANAGER_URL}`;
 const EVALUATOR_URL = `${process.env.VUE_APP_EVALUATOR_URL}`;
 const DATASTORES_URL = `${process.env.VUE_APP_DATASTORES_URL}`;
 const MODEL_MANAGER_URL = `${process.env.VUE_APP_MODEL_MANAGER_URL}`
-const LLM_MODELS_URL = MODEL_MANAGER_URL.replace('/models', '');
+const BASE_SQUARE_URL = MODEL_MANAGER_URL.replace('/models', '');
 
 /**
  * Get a list of available skill types.
@@ -263,15 +263,15 @@ export function getLocalLLMs() {
 export function generateText(params, streaming) {
     let url;
     if(!streaming){
-        url = `${LLM_MODELS_URL}/${params.model_identifier}/worker_generate`
+        url = `${BASE_SQUARE_URL}/${params.model_identifier}/worker_generate`
         const response = axios.post(url, params, {
             headers:{
                 'Content-Type': 'application/json'
             }
         });
         return response;
     }else{
-        url = `${LLM_MODELS_URL}/${params.model_identifier}/worker_generate_stream`
+        url = `${BASE_SQUARE_URL}/${params.model_identifier}/worker_generate_stream`
         const response = fetch(url, {
             method: 'POST',
             headers: {
@@ -288,10 +288,42 @@ export function getAlternatives(text){
     const params = {
         text: text
     }
-    const response = axios.post(`${LLM_MODELS_URL}/sensitivity/generate_alternatives`, params, {
+    const response = axios.post(`${BASE_SQUARE_URL}/sensitivity/generate_alternatives`, params, {
         headers:{
             'Content-Type': 'application/json'
         }
     });
     return response;
-}
+}
+
+export function getReplicateModels(){
+    return axios.get(`${BASE_SQUARE_URL}/replicate/models`, {});
+}
+
+export function generateChatStreamReplicate(params, token) {
+    let url = `${BASE_SQUARE_URL}/replicate/generate_chat_stream`
+    const response = axios.post(url, params, {
+        headers:{
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer ' + token
+        }
+    });
+    return response;
+}
+
+
+export function generateCompletionStreamReplicate(params, token) {
+    let url = `${BASE_SQUARE_URL}/replicate/generate_completion_stream`
+    const response = axios.post(url, params, {
+        headers:{
+            'Content-Type': 'application/json',
+            'Authorization': 'Bearer ' + token
+        }
+    });
+    return response;
+}
+
+
+export function replicateModelHealthCheck(model_id){
+    return axios.get(`${BASE_SQUARE_URL}/replicate/${model_id}/status`);
+}