HoverLingo / app.py
Gabriel382's picture
setting model manually
b1ed57d
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
@app.post("/translate")
def translate_text(request: TextRequest):
translated = translator(request.text, src_lang="en", tgt_lang="fr")[0]['translation_text']
return {"translated_text": translated}
@app.get("/")
def home():
return {"message": "Translation server is running!"}