amine_dubs commited on
Commit
6f3aee6
·
1 Parent(s): 068c749
Files changed (1) hide show
  1. backend/main.py +14 -7
backend/main.py CHANGED
@@ -51,11 +51,12 @@ os.environ['XDG_CACHE_HOME'] = '/tmp/cache'
51
  # --- Global model and tokenizer variables ---
52
  translator = None
53
  tokenizer = None
 
54
 
55
  # --- Model initialization function ---
56
  def initialize_model():
57
  """Initialize the translation model and tokenizer."""
58
- global translator, tokenizer
59
 
60
  try:
61
  print("Initializing model and tokenizer...")
@@ -69,16 +70,22 @@ def initialize_model():
69
  cache_dir="/tmp/transformers_cache"
70
  )
71
 
72
- # Create a pipeline for text2text generation
73
- # Important: Add from_tf=True to load TensorFlow weights
 
 
 
 
 
 
 
 
74
  translator = pipeline(
75
  "text2text-generation",
76
- model=model_name,
77
  tokenizer=tokenizer,
78
  device=-1, # Use CPU for compatibility (-1) or GPU if available (0)
79
- cache_dir="/tmp/transformers_cache",
80
- max_length=512,
81
- model_kwargs={"from_tf": True} # This is the key fix
82
  )
83
 
84
  print(f"Model {model_name} successfully initialized")
 
51
  # --- Global model and tokenizer variables ---
52
  translator = None
53
  tokenizer = None
54
+ model = None
55
 
56
  # --- Model initialization function ---
57
  def initialize_model():
58
  """Initialize the translation model and tokenizer."""
59
+ global translator, tokenizer, model
60
 
61
  try:
62
  print("Initializing model and tokenizer...")
 
70
  cache_dir="/tmp/transformers_cache"
71
  )
72
 
73
+ # Load the model explicitly with from_tf=True
74
+ print("Loading model with from_tf=True...")
75
+ model = AutoModelForSeq2SeqLM.from_pretrained(
76
+ model_name,
77
+ from_tf=True, # Explicitly set from_tf=True
78
+ cache_dir="/tmp/transformers_cache"
79
+ )
80
+
81
+ # Create a pipeline with the loaded model and tokenizer
82
+ print("Creating pipeline with pre-loaded model...")
83
  translator = pipeline(
84
  "text2text-generation",
85
+ model=model, # Use the model we loaded with from_tf=True
86
  tokenizer=tokenizer,
87
  device=-1, # Use CPU for compatibility (-1) or GPU if available (0)
88
+ max_length=512
 
 
89
  )
90
 
91
  print(f"Model {model_name} successfully initialized")