ClergeF's picture
Update app.py
ec4ae5f verified
raw
history blame contribute delete
645 Bytes
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
app = FastAPI()
# Load model
tokenizer = AutoTokenizer.from_pretrained("ClergeF/ms-hardwick-translator", use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained("ClergeF/ms-hardwick-translator")
pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
class InputText(BaseModel):
text: str
@app.post("/translate")
def translate(input: InputText):
output = pipe(input.text, max_length=512, do_sample=False)
return {"input": input.text, "output": output[0]["generated_text"]}