TiberiuCristianLeon commited on
Commit
f7cc9ad
·
verified ·
1 Parent(s): 0893423

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -2
app.py CHANGED
@@ -1,9 +1,12 @@
1
  from fastapi import FastAPI
2
  import src.paraphrase.Paraphrase as Paraphrase
3
  import src.translate.Translate as Translate
 
4
 
5
  app = FastAPI(docs_url="/")
6
- MTMODELS = {'enro': 'BlackKakapo/opus-mt-en-ro', 'roen': 'BlackKakapo/opus-mt-ro-en'}
 
 
7
 
8
  @app.get("/")
9
  def index():
@@ -15,7 +18,7 @@ def paraphrase(text: str, model: str):
15
  return {"input": text, "translation": resultValue, "exception": exception}
16
 
17
  @app.get("/translate")
18
- def translate(text: str, model: str):
19
  # resultValue, exception = Translate.paraphraseTranslateMethod(text, model)
20
  resultValue, exception = Translate.gemma_direct(text, model)
21
  return {"input": text, "translation": resultValue, "exception": exception}
 
1
  from fastapi import FastAPI
2
  import src.paraphrase.Paraphrase as Paraphrase
3
  import src.translate.Translate as Translate
4
+ from typing import Optional
5
 
6
  app = FastAPI(docs_url="/")
7
+ MTMODELS = {'enro': 'BlackKakapo/opus-mt-en-ro',
8
+ 'roen': 'BlackKakapo/opus-mt-ro-en',
9
+ 'gemma': 'Gargaz/gemma-2b-romanian-better'}
10
 
11
  @app.get("/")
12
  def index():
 
18
  return {"input": text, "translation": resultValue, "exception": exception}
19
 
20
  @app.get("/translate")
21
+ def translate(text: str, model: Optional[str] = MTMODELS['gemma']):
22
  # resultValue, exception = Translate.paraphraseTranslateMethod(text, model)
23
  resultValue, exception = Translate.gemma_direct(text, model)
24
  return {"input": text, "translation": resultValue, "exception": exception}