Spaces:
Sleeping
Sleeping
| import os | |
| from fastapi import FastAPI | |
| from pydantic import BaseModel | |
| from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer | |
| # β Set cache directory inside Hugging Face's allowed `/tmp/` folder | |
| CACHE_DIR = "/tmp/huggingface_cache" | |
| os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR | |
| os.environ["HF_HOME"] = CACHE_DIR | |
| os.environ["HF_HUB_DISABLE_SYMLINKS"] = "1" # Fixes permission issue | |
| # β Ensure cache directory exists | |
| os.makedirs(CACHE_DIR, exist_ok=True) | |
| # β Load the model | |
| MODEL_NAME = "facebook/m2m100_418M" | |
| print("Downloading model...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR) | |
| translator = pipeline("translation", model=model, tokenizer=tokenizer) | |
| print("Model loaded successfully!") | |
| # β FastAPI App | |
| app = FastAPI() | |
| class TextRequest(BaseModel): | |
| text: str | |
| def translate_text(request: TextRequest): | |
| translated = translator(request.text, src_lang="en", tgt_lang="fr")[0]['translation_text'] | |
| return {"translated_text": translated} | |
| def home(): | |
| return {"message": "Translation server is running!"} | |