from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse from pydantic import BaseModel from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import torch import asyncio import httpx import time from concurrent.futures import ThreadPoolExecutor import hashlib from collections import defaultdict app = FastAPI(title="Distributed M2M100 API") # ===================================================== # CORS # ===================================================== app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ===================================================== # MODEL # ===================================================== MODEL_NAME = "facebook/m2m100_418M" print("Loading model...") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) model.eval() torch.set_num_threads(2) torch.set_num_interop_threads(1) print("Model loaded!") # ===================================================== # CONFIG # ===================================================== MAX_ACTIVE_REQUESTS = 3 MAX_QUEUE_SIZE = 3 MAX_CHARACTERS = 3000 # Character limit without spaces per user USERNAME = "velosts" SPACE_PREFIX = "trans" DISCOVERY_CHUNK_SIZE = 50 CACHE_TTL = 86400 # Cache time-to-live: 1 day in seconds (24 * 60 * 60) MAX_CACHE_SIZE = 10000 # Maximum number of cached translations MAX_RPS_LANGUAGES = 20000 # Maximum requests per second for /languages endpoint # ===================================================== # GLOBALS # ===================================================== active_requests = 0 request_queue = asyncio.Queue(maxsize=MAX_QUEUE_SIZE) executor = ThreadPoolExecutor(max_workers=3) lock = asyncio.Lock() # 🔥 cooldown map (NEW) cooldown = {} # 🔥 Translation cache with TTL (NEW) translation_cache = {} # 🔥 RPS tracking for /languages endpoint (NEW) languages_request_times = [] languages_lock = asyncio.Lock() # ===================================================== # LANGUAGES # ===================================================== SUPPORTED_LANGS = { "af": "Afrikaans", "am": "Amharic", "ar": "Arabic", "ast": "Asturian", "az": "Azerbaijani", "ba": "Bashkir", "be": "Belarusian", "bg": "Bulgarian", "bn": "Bengali", "br": "Breton", "bs": "Bosnian", "ca": "Catalan", "ceb": "Cebuano", "cs": "Czech", "cy": "Welsh", "da": "Danish", "de": "German", "el": "Greek", "en": "English", "es": "Spanish", "et": "Estonian", "fa": "Persian", "ff": "Fulah", "fi": "Finnish", "fr": "French", "fy": "Western Frisian", "ga": "Irish", "gd": "Scottish Gaelic", "gl": "Galician", "gu": "Gujarati", "ha": "Hausa", "he": "Hebrew", "hi": "Hindi", "hr": "Croatian", "ht": "Haitian Creole", "hu": "Hungarian", "hy": "Armenian", "id": "Indonesian", "ig": "Igbo", "ilo": "Ilocano", "is": "Icelandic", "it": "Italian", "ja": "Japanese", "jv": "Javanese", "ka": "Georgian", "kk": "Kazakh", "km": "Khmer", "kn": "Kannada", "ko": "Korean", "lb": "Luxembourgish", "lg": "Ganda", "ln": "Lingala", "lo": "Lao", "lt": "Lithuanian", "lv": "Latvian", "mg": "Malagasy", "mk": "Macedonian", "ml": "Malayalam", "mn": "Mongolian", "mr": "Marathi", "ms": "Malay", "my": "Myanmar", "ne": "Nepali", "nl": "Dutch", "no": "Norwegian", "ns": "Northern Sotho", "oc": "Occitan", "or": "Odia", "pa": "Punjabi", "pl": "Polish", "ps": "Pashto", "pt": "Portuguese", "ro": "Romanian", "ru": "Russian", "sd": "Sindhi", "si": "Sinhala", "sk": "Slovak", "sl": "Slovenian", "so": "Somali", "sq": "Albanian", "sr": "Serbian", "ss": "Swati", "su": "Sundanese", "sv": "Swedish", "sw": "Swahili", "ta": "Tamil", "th": "Thai", "tl": "Tagalog", "tn": "Tswana", "tr": "Turkish", "uk": "Ukrainian", "ur": "Urdu", "uz": "Uzbek", "vi": "Vietnamese", "wo": "Wolof", "xh": "Xhosa", "yi": "Yiddish", "yo": "Yoruba", "zh": "Chinese", "zu": "Zulu" } # ===================================================== # REQUEST MODEL # ===================================================== class TranslateRequest(BaseModel): text: str source_lang: str target_lang: str # ===================================================== # TRANSLATION # ===================================================== def blocking_translate(text, source_lang, target_lang): tokenizer.src_lang = source_lang encoded = tokenizer( text, return_tensors="pt" ).to(device) generated_tokens = model.generate( **encoded, forced_bos_token_id=tokenizer.get_lang_id(target_lang), max_length=1024, num_beams=1 ) return tokenizer.batch_decode( generated_tokens, skip_special_tokens=True )[0] # ===================================================== # CACHE HELPERS (NEW) # ===================================================== def get_cache_key(text, source_lang, target_lang): """Generate a unique cache key for a translation request""" raw_key = f"{text}|{source_lang}|{target_lang}" return hashlib.md5(raw_key.encode()).hexdigest() def get_from_cache(text, source_lang, target_lang): """Retrieve translation from cache if valid""" cache_key = get_cache_key(text, source_lang, target_lang) if cache_key in translation_cache: cached_data = translation_cache[cache_key] if time.time() - cached_data["timestamp"] < CACHE_TTL: return cached_data["translated_text"] else: # Remove expired cache entry del translation_cache[cache_key] return None def add_to_cache(text, source_lang, target_lang, translated_text): """Add translation to cache with timestamp. Override oldest entry if cache is full.""" cache_key = get_cache_key(text, source_lang, target_lang) # If key already exists, update it if cache_key in translation_cache: translation_cache[cache_key] = { "translated_text": translated_text, "timestamp": time.time() } return # If cache is full, find and remove the oldest entry if len(translation_cache) >= MAX_CACHE_SIZE: oldest_key = min( translation_cache.keys(), key=lambda k: translation_cache[k]["timestamp"] ) del translation_cache[oldest_key] # Add new entry translation_cache[cache_key] = { "translated_text": translated_text, "timestamp": time.time() } # ===================================================== # COOLDOWN HELPERS # ===================================================== def is_blocked(url): return cooldown.get(url, 0) > time.time() # ===================================================== # CHECK SINGLE SPACE # ===================================================== async def check_space(client, i): space_name = SPACE_PREFIX if i == 0 else f"{SPACE_PREFIX}{i}" url = f"https://{USERNAME}-{space_name}.hf.space" if is_blocked(url): return { "exists": True, "space_name": space_name, "url": url, "status_data": { "current_space_status": "cooldown" } } try: response = await client.get(f"{url}/status") if response.status_code == 200: try: return { "exists": True, "space_name": space_name, "url": url, "status_data": response.json() } except: return { "exists": True, "space_name": space_name, "url": url, "status_data": None } elif response.status_code == 404: return { "exists": False, "space_name": space_name, "url": url } elif response.status_code == 429: cooldown[url] = time.time() + 1 return { "exists": True, "space_name": space_name, "url": url, "status_data": { "current_space_status": "rate_limited" } } else: return { "exists": True, "space_name": space_name, "url": url, "status_data": { "current_space_status": "error" } } except: return { "exists": False, "space_name": space_name, "url": url } # ===================================================== # DISCOVERY # ===================================================== async def discover_spaces_parallel(): existing_spaces = [] empty_space_url = None timeout = httpx.Timeout(5.0) async with httpx.AsyncClient(timeout=timeout) as client: start = 0 while True: tasks = [ check_space(client, i) for i in range(start, start + DISCOVERY_CHUNK_SIZE) ] results = await asyncio.gather(*tasks) found_empty = False for result in results: if result["exists"]: status = (result.get("status_data") or {}).get("current_space_status") if status in ["rate_limited", "cooldown"]: continue existing_spaces.append({ "space_name": result["space_name"], "url": result["url"], "status": result.get("status_data") }) else: empty_space_url = result["url"] found_empty = True break if found_empty: break start += DISCOVERY_CHUNK_SIZE return { "existing_spaces": existing_spaces, "empty_space_url": empty_space_url } # ===================================================== # STATUS # ===================================================== @app.get("/status") async def status(): queue_size = request_queue.qsize() full = ( active_requests >= MAX_ACTIVE_REQUESTS and queue_size >= MAX_QUEUE_SIZE ) discovered = await discover_spaces_parallel() return { "current_space_status": "full" if full else "available", "active_requests": active_requests, "queue_size": queue_size, "empty_space_url": discovered["empty_space_url"], "total_spaces": len(discovered["existing_spaces"]), "cached_translations": len(translation_cache) } # ===================================================== # LANGUAGES WITH RPS LIMITING (UPDATED) # ===================================================== @app.get("/languages") async def languages(request: Request): async with languages_lock: current_time = time.time() # Remove requests older than 1 second global languages_request_times languages_request_times = [ t for t in languages_request_times if current_time - t < 1.0 ] # Check if we've exceeded the RPS limit if len(languages_request_times) >= MAX_RPS_LANGUAGES: # Discover empty space URL for redirection discovered = await discover_spaces_parallel() return JSONResponse( status_code=429, content={ "status": "space_full", "message": f"Rate limit exceeded. Maximum {MAX_RPS_LANGUAGES} requests per second allowed.", "empty_space_url": discovered["empty_space_url"] } ) # Add current request timestamp languages_request_times.append(current_time) return SUPPORTED_LANGS # ===================================================== # TRANSLATE # ===================================================== @app.post("/translate") async def translate(req: TranslateRequest): global active_requests # Character limit check (without spaces) char_count_without_spaces = len(req.text.replace(" ", "")) if char_count_without_spaces > MAX_CHARACTERS: return { "status": "character_limit_exceeded", "message": f"Text exceeds maximum character limit (without spaces). Maximum allowed: {MAX_CHARACTERS} characters, received: {char_count_without_spaces} characters", "max_characters": MAX_CHARACTERS, "received_characters": char_count_without_spaces } # Check cache first cached_result = get_from_cache(req.text, req.source_lang, req.target_lang) if cached_result is not None: return { "status": "success", "translated_text": cached_result, "from_cache": True } queue_size = request_queue.qsize() if active_requests >= MAX_ACTIVE_REQUESTS and queue_size >= MAX_QUEUE_SIZE: discovered = await discover_spaces_parallel() return { "status": "space_full", "empty_space_url": discovered["empty_space_url"] } await request_queue.put("req") async with lock: active_requests += 1 try: loop = asyncio.get_running_loop() result = await loop.run_in_executor( executor, blocking_translate, req.text, req.source_lang, req.target_lang ) # Store result in cache add_to_cache(req.text, req.source_lang, req.target_lang, result) return { "status": "success", "translated_text": result, "from_cache": False } finally: async with lock: active_requests -= 1 await request_queue.get() request_queue.task_done() # ===================================================== # HOME # ===================================================== @app.get("/") def home(): return {"status": "running"}