amine_dubs commited on
Commit
20ee4d2
·
1 Parent(s): 9f23223

Fix: Explicitly pass cache_dir to from_pretrained

Browse files
Files changed (1) hide show
  1. backend/main.py +5 -4
backend/main.py CHANGED
@@ -33,16 +33,17 @@ templates = Jinja2Templates(directory=TEMPLATE_DIR)
33
 
34
  # Define model name
35
  MODEL_NAME = "Helsinki-NLP/opus-mt-en-ar"
 
36
  translator = None # Initialize translator as None
37
 
38
  try:
39
  print("--- Loading Model ---") # Add a clear marker
40
  print(f"Loading tokenizer for {MODEL_NAME} using MarianTokenizer...")
41
- # Use MarianTokenizer directly
42
- tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
43
  print(f"Loading model for {MODEL_NAME} using MarianMTModel...")
44
- # Use MarianMTModel directly
45
- model = MarianMTModel.from_pretrained(MODEL_NAME)
46
  print(f"Initializing translation pipeline for {MODEL_NAME}...")
47
  # Pass the loaded objects to the pipeline
48
  translator = pipeline("translation", model=model, tokenizer=tokenizer)
 
33
 
34
  # Define model name
35
  MODEL_NAME = "Helsinki-NLP/opus-mt-en-ar"
36
+ CACHE_DIR = "/app/.cache" # Explicitly define cache directory
37
  translator = None # Initialize translator as None
38
 
39
  try:
40
  print("--- Loading Model ---") # Add a clear marker
41
  print(f"Loading tokenizer for {MODEL_NAME} using MarianTokenizer...")
42
+ # Use MarianTokenizer directly and specify cache_dir
43
+ tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
44
  print(f"Loading model for {MODEL_NAME} using MarianMTModel...")
45
+ # Use MarianMTModel directly and specify cache_dir
46
+ model = MarianMTModel.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
47
  print(f"Initializing translation pipeline for {MODEL_NAME}...")
48
  # Pass the loaded objects to the pipeline
49
  translator = pipeline("translation", model=model, tokenizer=tokenizer)