FastAPIMT / app.py
TiberiuCristianLeon's picture
Update app.py (#9)
9f0a178 verified
raw
history blame contribute delete
929 Bytes
from fastapi import FastAPI
import src.Paraphrase as Paraphrase
import src.Translate as Translate
from typing import Optional
app = FastAPI(docs_url="/")
MTMODELS = {'enro': 'BlackKakapo/opus-mt-en-ro',
'roen': 'BlackKakapo/opus-mt-ro-en',
'gemma': 'Gargaz/gemma-2b-romanian-better'}
@app.get("/")
def index():
return {'endpoints': ['/paraphrase', '/translate'], 'mtmodels': MTMODELS}
@app.get("/paraphrase")
def paraphrase(text: str, model: str):
resultValue, exception = Paraphrase.paraphraseParaphraseMethod(text, model)
return {"input": text, "translation": resultValue, "exception": exception}
@app.get("/translate")
def translate(text: str, model: Optional[str] = MTMODELS['gemma']):
# resultValue, exception = Translate.paraphraseTranslateMethod(text, model)
resultValue = Translate.gemma_direct(text, model)
return {"input": text, "result": resultValue, "model": model}