Spaces:
Sleeping
Sleeping
File size: 3,481 Bytes
7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 42635f8 d4a6660 42635f8 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 7097550 d4a6660 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
from IndicTransToolkit.processor import IndicProcessor
DEVICE = "cpu" # HF free tier
# -------------------------------
# Models
# -------------------------------
INDIC_EN_MODEL = "ai4bharat/indictrans2-indic-en-1B"
EN_INDIC_MODEL = "ai4bharat/indictrans2-en-indic-1B"
# -------------------------------
# Load Indic → English
# -------------------------------
indic_en_tokenizer = AutoTokenizer.from_pretrained(
INDIC_EN_MODEL,
trust_remote_code=True
)
indic_en_model = AutoModelForSeq2SeqLM.from_pretrained(
INDIC_EN_MODEL,
trust_remote_code=True
).to(DEVICE)
indic_en_model.eval()
# -------------------------------
# Load English → Indic
# -------------------------------
en_indic_tokenizer = AutoTokenizer.from_pretrained(
EN_INDIC_MODEL,
trust_remote_code=True
)
en_indic_model = AutoModelForSeq2SeqLM.from_pretrained(
EN_INDIC_MODEL,
trust_remote_code=True
).to(DEVICE)
en_indic_model.eval()
# -------------------------------
# Processor
# -------------------------------
ip = IndicProcessor(inference=True)
# -------------------------------
# FastAPI app
# -------------------------------
app = FastAPI(
title="Indic ↔ English Translation API",
docs_url="/docs"
)
# -------------------------------
# Schemas
# -------------------------------
class IndicToEnRequest(BaseModel):
text: str
src_lang: str # e.g. kan_Knda, hin_Deva
class EnToIndicRequest(BaseModel):
text: str
tgt_lang: str # e.g. kan_Knda, hin_Deva
# -------------------------------
# Root
# -------------------------------
@app.get("/")
def root():
return {
"status": "ok",
"endpoints": {
"indic_to_en": "/translate/indic-to-en",
"en_to_indic": "/translate/en-to-indic"
}
}
# -------------------------------
# Indic → English
# -------------------------------
@app.post("/translate/indic-to-en")
def indic_to_en(req: IndicToEnRequest):
batch = ip.preprocess_batch(
[req.text],
src_lang=req.src_lang,
tgt_lang="eng_Latn"
)
inputs = indic_en_tokenizer(
batch, return_tensors="pt", padding=True
)
with torch.no_grad():
outputs = indic_en_model.generate(
**inputs,
max_length=128,
num_beams=3,
use_cache=False
)
translation = indic_en_tokenizer.batch_decode(
outputs, skip_special_tokens=True
)[0]
translation = ip.postprocess_batch(
[translation], "eng_Latn"
)[0]
return {"translation": translation}
# -------------------------------
# English → Indic
# -------------------------------
@app.post("/translate/en-to-indic")
def en_to_indic(req: EnToIndicRequest):
batch = ip.preprocess_batch(
[req.text],
src_lang="eng_Latn",
tgt_lang=req.tgt_lang
)
inputs = en_indic_tokenizer(
batch, return_tensors="pt", padding=True
)
with torch.no_grad():
outputs = en_indic_model.generate(
**inputs,
max_length=128,
num_beams=3,
use_cache=False
)
translation = en_indic_tokenizer.batch_decode(
outputs, skip_special_tokens=True
)[0]
translation = ip.postprocess_batch(
[translation], req.tgt_lang
)[0]
return {"translation": translation}
|