start teste

Dockerfile

@@ -0,0 +1,11 @@
+FROM tiangolo/uvicorn-gunicorn-fastapi:python3.9
+
+WORKDIR /app
+
+COPY requirements.txt .
+
+RUN pip install -r requirements.txt
+
+COPY . .
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "80"]

main.py

@@ -0,0 +1,50 @@
+import io
+import time
+from typing import List, Literal
+from fastapi import FastAPI
+from pydantic import BaseModel
+from enum import Enum
+from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+import torch
+
+app = FastAPI()
+device = torch.device("cpu")
+
+
+class TranslationRequest(BaseModel):
+    user_input: str
+    source_lang: str
+    target_lang: str
+
+
+def load_model(pretrained_model: str = "facebook/m2m100_1.2B", cache_dir: str = "models/"):
+    tokenizer = M2M100Tokenizer.from_pretrained(pretrained_model, cache_dir=cache_dir)
+    model = M2M100ForConditionalGeneration.from_pretrained(pretrained_model, cache_dir=cache_dir).to(device)
+    model.eval()
+    return tokenizer, model
+
+
+@app.post("/translate")
+async def translate(request: TranslationRequest):
+    time_start = time.time()
+    tokenizer, model = load_model()
+    src_lang = request.source_lang
+    trg_lang = request.target_lang
+    tokenizer.src_lang = src_lang
+    with torch.no_grad():
+        encoded_input = tokenizer(request.user_input, return_tensors="pt").to(device)
+        generated_tokens = model.generate(
+            **encoded_input, forced_bos_token_id=tokenizer.get_lang_id(trg_lang)
+        )
+        translated_text = tokenizer.batch_decode(
+            generated_tokens, skip_special_tokens=True
+        )[0]
+    time_end = time.time()
+    response = {"translation": translated_text, "computation_time": round((time_end - time_start), 3)}
+    return response
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -0,0 +1,7 @@
+wheel
+fastapi[all]
+gunicorn
+streamlit
+torch
+transformers
+transformers[sentencepiece]