mirror of
https://github.com/vale981/ray
synced 2025-03-08 19:41:38 -05:00
42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
from fastapi import FastAPI
|
|
from transformers import pipeline
|
|
|
|
# TODO: actually use this to predict something
|
|
import ray
|
|
from ray import serve
|
|
|
|
app = FastAPI()
|
|
|
|
|
|
# Define our deployment.
|
|
@serve.deployment(num_replicas=2)
|
|
class GPT2:
|
|
def __init__(self):
|
|
self.nlp_model = pipeline("text-generation", model="gpt2")
|
|
|
|
async def predict(self, query: str):
|
|
return self.nlp_model(query, max_length=50)
|
|
|
|
async def __call__(self, request):
|
|
return self.predict(await request.body())
|
|
|
|
|
|
@app.on_event("startup") # Code to be run when the server starts.
|
|
async def startup_event():
|
|
ray.init(address="auto") # Connect to the running Ray cluster.
|
|
serve.start(http_host=None) # Start the Ray Serve instance.
|
|
|
|
# Deploy our GPT2 Deployment.
|
|
GPT2.deploy()
|
|
|
|
|
|
@app.get("/generate")
|
|
async def generate(query: str):
|
|
# Get a handle to our deployment so we can query it in Python.
|
|
handle = GPT2.get_handle()
|
|
return await handle.predict.remote(query)
|
|
|
|
|
|
@app.on_event("shutdown") # Code to be run when the server shuts down.
|
|
async def shutdown_event():
|
|
serve.shutdown() # Shut down Ray Serve.
|