Expose LLM API with FastAPI

Dockerfile

@@ -0,0 +1,16 @@
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+WORKDIR /app
+
+COPY --chown=user requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+
+USER user
+
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -5,6 +5,11 @@ colorFrom: pink
 colorTo: blue
 sdk: docker
 pinned: false
+app_port: 7860
 ---
 
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+
+

main.py

@@ -0,0 +1,20 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import pipeline
+
+app = FastAPI()
+
+pipe = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-small"
+)
+
+class Prompt(BaseModel):
+    input: str
+
+@app.post("/generate")
+def generate(prompt: Prompt):
+    output = pipe(prompt.input)
+    return {
+        "output": output[0]["generated_text"]
+    }

requirements.txt

@@ -0,0 +1,6 @@
+fastapi==0.74.*
+requests==2.27.*
+sentencepiece==0.1.*
+torch==1.11.*
+transformers==4.*
+uvicorn[standard]==0.17.*