fastapi backend

DOCKERFILE

@@ -0,0 +1,12 @@
+FROM nvidia/cuda:12.1.0-base-ubuntu22.04
+
+RUN apt-get update && apt-get install -y python3 python3-pip
+
+WORKDIR /app
+
+COPY requirements.txt .
+RUN pip3 install -r requirements.txt
+
+COPY app.py .
+
+CMD ["python3", "app.py"]

app.py

@@ -1,7 +1,28 @@
 from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from vllm import LLM, SamplingParams
 
 app = FastAPI()
 
-@app.get("/")
-def greet_json():
-    return {"Hello": "World!"}
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+llama = LLM(
+    model="meta-llama/Llama-3.2-1B-Instruct",
+    dtype="half",
+    gpu_memory_utilization=0.9
+)
+
+@app.post("/generate")
+def generate(prompt: str):
+    params = SamplingParams(temperature=0.7, max_tokens=100)
+    output = llama.generate([prompt], params)
+    return {"text": output[0].outputs[0].text}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -1,2 +1,7 @@
 fastapi
-uvicorn[standard]
+uvicorn
+vllm
+transformers
+torch
+pillow
+faster-whisper