omkar56 commited on
Commit
2d14bdf
·
1 Parent(s): d58981a

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +18 -18
main.py CHANGED
@@ -40,28 +40,28 @@ async def ocr(
40
  return jsonable_encoder({"text": text})
41
 
42
 
43
- @app.post("/api/translate", response_model=dict)
44
- async def translate(
45
- api_key: str = Depends(get_api_key),
46
- text: str = Body(...),
47
- src: str = "en",
48
- trg: str = "zh",
49
- ):
50
  # if api_key != API_KEY:
51
  # return {"error": "Invalid API key"}, 401
52
 
53
- tokenizer, model = get_model(src, trg)
54
 
55
- translated_text = ""
56
- for sentence in sent_tokenize(text):
57
- translated_sub = model.generate(**tokenizer(sentence, return_tensors="pt"))[0]
58
- translated_text += tokenizer.decode(translated_sub, skip_special_tokens=True) + "\n"
59
 
60
- return jsonable_encoder({"translated_text": translated_text})
61
 
62
 
63
- def get_model(src: str, trg: str):
64
- model_name = f"Helsinki-NLP/opus-mt-{src}-{trg}"
65
- tokenizer = MarianTokenizer.from_pretrained(model_name)
66
- model = MarianMTModel.from_pretrained(model_name)
67
- return tokenizer, model
 
40
  return jsonable_encoder({"text": text})
41
 
42
 
43
+ #@app.post("/api/translate", response_model=dict)
44
+ #async def translate(
45
+ #api_key: str = Depends(get_api_key),
46
+ #text: str = Body(...),
47
+ #src: str = "en",
48
+ #trg: str = "zh",
49
+ #):
50
  # if api_key != API_KEY:
51
  # return {"error": "Invalid API key"}, 401
52
 
53
+ # tokenizer, model = get_model(src, trg)
54
 
55
+ # translated_text = ""
56
+ # for sentence in sent_tokenize(text):
57
+ # translated_sub = model.generate(**tokenizer(sentence, return_tensors="pt"))[0]
58
+ # translated_text += tokenizer.decode(translated_sub, skip_special_tokens=True) + "\n"
59
 
60
+ # return jsonable_encoder({"translated_text": translated_text})
61
 
62
 
63
+ #def get_model(src: str, trg: str):
64
+ # model_name = f"Helsinki-NLP/opus-mt-{src}-{trg}"
65
+ #tokenizer = MarianTokenizer.from_pretrained(model_name)
66
+ #model = MarianMTModel.from_pretrained(model_name)
67
+ #return tokenizer, model