Spaces:
Running
Running
Darshan
commited on
Commit
•
d39f3fd
1
Parent(s):
81d5bff
add apis
Browse files- Dockerfile +14 -0
- app.py +83 -0
Dockerfile
ADDED
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Use the official Python 3.10.9 image
|
2 |
+
FROM python:3.10.9
|
3 |
+
|
4 |
+
# Copy the current directory contents into the container at .
|
5 |
+
COPY . .
|
6 |
+
|
7 |
+
# Set the working directory to /
|
8 |
+
WORKDIR /
|
9 |
+
|
10 |
+
# Install requirements.txt
|
11 |
+
RUN pip install --no-cache-dir --upgrade -r /requirements.txt
|
12 |
+
|
13 |
+
# Start the FastAPI app on port 7860, the default port expected by Spaces
|
14 |
+
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
|
app.py
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from fastapi import FastAPI
|
2 |
+
from typing import List
|
3 |
+
import torch
|
4 |
+
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
|
5 |
+
from IndicTransToolkit import IndicProcessor
|
6 |
+
from fastapi.middleware.cors import CORSMiddleware
|
7 |
+
|
8 |
+
# Initialize FastAPI
|
9 |
+
app = FastAPI()
|
10 |
+
|
11 |
+
# Add CORS middleware
|
12 |
+
app.add_middleware(
|
13 |
+
CORSMiddleware,
|
14 |
+
allow_origins=["*"],
|
15 |
+
allow_credentials=True,
|
16 |
+
allow_methods=["*"],
|
17 |
+
allow_headers=["*"],
|
18 |
+
)
|
19 |
+
|
20 |
+
# Initialize models and processors
|
21 |
+
model = AutoModelForSeq2SeqLM.from_pretrained(
|
22 |
+
"ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
|
23 |
+
)
|
24 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
25 |
+
"ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True
|
26 |
+
)
|
27 |
+
ip = IndicProcessor(inference=True)
|
28 |
+
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
|
29 |
+
model = model.to(DEVICE)
|
30 |
+
|
31 |
+
|
32 |
+
def translate_text(sentences: List[str], target_lang: str):
|
33 |
+
try:
|
34 |
+
src_lang = "eng_Latn"
|
35 |
+
batch = ip.preprocess_batch(sentences, src_lang=src_lang, tgt_lang=target_lang)
|
36 |
+
inputs = tokenizer(
|
37 |
+
batch,
|
38 |
+
truncation=True,
|
39 |
+
padding="longest",
|
40 |
+
return_tensors="pt",
|
41 |
+
return_attention_mask=True,
|
42 |
+
).to(DEVICE)
|
43 |
+
|
44 |
+
with torch.no_grad():
|
45 |
+
generated_tokens = model.generate(
|
46 |
+
**inputs,
|
47 |
+
use_cache=True,
|
48 |
+
min_length=0,
|
49 |
+
max_length=256,
|
50 |
+
num_beams=5,
|
51 |
+
num_return_sequences=1,
|
52 |
+
)
|
53 |
+
|
54 |
+
with tokenizer.as_target_tokenizer():
|
55 |
+
generated_tokens = tokenizer.batch_decode(
|
56 |
+
generated_tokens.detach().cpu().tolist(),
|
57 |
+
skip_special_tokens=True,
|
58 |
+
clean_up_tokenization_spaces=True,
|
59 |
+
)
|
60 |
+
|
61 |
+
translations = ip.postprocess_batch(generated_tokens, lang=target_lang)
|
62 |
+
return {
|
63 |
+
"translations": translations,
|
64 |
+
"source_language": src_lang,
|
65 |
+
"target_language": target_lang,
|
66 |
+
}
|
67 |
+
except Exception as e:
|
68 |
+
raise Exception(f"Translation failed: {str(e)}")
|
69 |
+
|
70 |
+
|
71 |
+
# FastAPI routes
|
72 |
+
@app.get("/health")
|
73 |
+
async def health_check():
|
74 |
+
return {"status": "healthy"}
|
75 |
+
|
76 |
+
|
77 |
+
@app.post("/translate")
|
78 |
+
async def translate_endpoint(sentences: List[str], target_lang: str):
|
79 |
+
try:
|
80 |
+
result = translate_text(sentences=sentences, target_lang=target_lang)
|
81 |
+
return result
|
82 |
+
except Exception as e:
|
83 |
+
raise HTTPException(status_code=500, detail=str(e))
|