Upload 3 files

Dockerfile.txt

@@ -0,0 +1,17 @@
+# Use the official Python image from Docker Hub
+FROM python:3.10-slim
+
+# Set the working directory
+WORKDIR /app
+
+# Copy the requirements file into the container
+COPY requirements.txt .
+
+# Install dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the current directory contents into the container at /app
+COPY . .
+
+# Command to run the FastAPI app
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,31 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
+
+app = FastAPI()
+
+# Load model once on startup
+MODEL_NAME = "facebook/nllb-200-1.3B"  # 3B version
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
+
+class TranslationRequest(BaseModel):
+    text: str
+    src_lang: str
+    tgt_lang: str
+
+@app.post("/translate")
+def translate(req: TranslationRequest):
+    inputs = tokenizer(
+        req.text,
+        return_tensors="pt",
+    ).to(model.device)
+
+    outputs = model.generate(
+        **inputs,
+        forced_bos_token_id=tokenizer.lang_code_to_id[req.tgt_lang],
+        max_length=512
+    )
+    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"translation": translation}

requirements.txt

@@ -0,0 +1,7 @@
+transformers
+torch
+sentencepiece
+tiktoken
+fastapi
+uvicorn
+protobuf