Darshan commited on
Commit
d39f3fd
1 Parent(s): 81d5bff
Files changed (2) hide show
  1. Dockerfile +14 -0
  2. 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))