Add workflow logic for weather forecast demo (#24)

embedding-server/Dockerfile

@@ -24,6 +24,7 @@ FROM python:3-slim AS output
 # following models have been tested to work with this image
 # "sentence-transformers/all-MiniLM-L6-v2,sentence-transformers/all-mpnet-base-v2,thenlper/gte-base,thenlper/gte-large,thenlper/gte-small"
 ENV MODELS="BAAI/bge-large-en-v1.5"
+ENV NER_MODELS="urchade/gliner_large-v2.1"
 
 COPY --from=builder /runtime /usr/local
 

embedding-server/app/install.py

@@ -1,3 +1,6 @@
-from load_transformers import load_transformers
+from load_models import load_transformers, load_ner_models
 
+print('installing transformers')
 load_transformers()
+print('installing ner models')
+load_ner_models()

embedding-server/app/load_models.py

@@ -1,5 +1,6 @@
 import os
 import sentence_transformers
+from gliner import GLiNER
 
 def load_transformers(models = os.getenv("MODELS", "sentence-transformers/all-MiniLM-L6-v2")):
     transformers = {}
@@ -8,3 +9,11 @@ def load_transformers(models = os.getenv("MODELS", "sentence-transformers/all-Mi
         transformers[model] = sentence_transformers.SentenceTransformer(model)
 
     return transformers
+
+def load_ner_models(models = os.getenv("NER_MODELS", "urchade/gliner_large-v2.1")):
+    ner_models = {}
+
+    for model in models.split(','):
+        ner_models[model] = GLiNER.from_pretrained(model)
+
+    return ner_models

embedding-server/app/main.py

@@ -1,8 +1,11 @@
+import random
 from fastapi import FastAPI, Response, HTTPException
 from pydantic import BaseModel
-from load_transformers import load_transformers
+from load_models import load_ner_models, load_transformers
+from datetime import date, timedelta
 
 transformers = load_transformers()
+ner_models = load_ner_models()
 
 app = FastAPI()
 
@@ -10,6 +13,12 @@ class EmbeddingRequest(BaseModel):
   input: str
   model: str
 
+@app.get("/healthz")
+async def healthz():
+    return {
+        "status": "ok"
+    }
+
 @app.get("/models")
 async def models():
     models = []
@@ -27,7 +36,7 @@ async def models():
 
 @app.post("/embeddings")
 async def embedding(req: EmbeddingRequest, res: Response):
-    if not req.model in transformers:
+    if req.model not in transformers:
         raise HTTPException(status_code=400, detail="unknown model: " + req.model)
 
     embeddings = transformers[req.model].encode([req.input])
@@ -51,3 +60,48 @@ async def embedding(req: EmbeddingRequest, res: Response):
         "object": "list",
         "usage": usage
     }
+
+class NERRequest(BaseModel):
+  input: str
+  labels: list[str]
+  model: str
+
+
+@app.post("/ner")
+async def ner(req: NERRequest, res: Response):
+    if req.model not in ner_models:
+        raise HTTPException(status_code=400, detail="unknown model: " + req.model)
+
+    model = ner_models[req.model]
+    entities = model.predict_entities(req.input, req.labels)
+
+    return {
+        "data": entities,
+        "model": req.model,
+        "object": "list",
+    }
+
+class WeatherRequest(BaseModel):
+  city: str
+
+
+@app.post("/weather")
+async def weather(req: WeatherRequest, res: Response):
+
+    weather_forecast = {
+        "city": req.city,
+        "temperature": [],
+        "unit": "F",
+    }
+    for i in range(7):
+       min_temp = random.randrange(50,90)
+       max_temp = random.randrange(min_temp+5, min_temp+20)
+       weather_forecast["temperature"].append({
+           "date": str(date.today() + timedelta(days=i)),
+           "temperature": {
+              "min": min_temp,
+              "max": max_temp
+           }
+       })
+
+    return weather_forecast

embedding-server/requirements.txt

@@ -3,3 +3,4 @@ fastapi
 sentence-transformers
 torch
 uvicorn
+gliner