DockerDemo / app.py
SGK070403's picture
Update app.py
54a0dd6 verified
Raw
History Blame Contribute Delete
520 Bytes
from fastapi import FastAPI
from transformers import pipeline
import torch
app = FastAPI()
pipe = pipeline("text2text-generation", model="lmsys/fastchat-t5-3b-v1.0")
messages = [
{"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
{"role": "user", "content": "Who are you?"},
]
@app.get("/")
def home():
return {"message":"Hello world"}
@app.get("/generate")
def generate(text:str):
outputs = pipe(text)
return{"output": outputs[0]["generated_text"]}