Erfan11 commited on
Commit
da9a7a0
1 Parent(s): caa9290

Update load_model.py

Browse files
Files changed (1) hide show
  1. load_model.py +3 -3
load_model.py CHANGED
@@ -8,13 +8,13 @@ load_dotenv()
8
  def load_model(model_name):
9
  try:
10
  # Load TensorFlow model from Hugging Face
11
- model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'), from_tf=True)
12
  except OSError:
13
  raise ValueError("Model loading failed.")
14
  return model
15
 
16
  def load_tokenizer(model_name):
17
- tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('API_KEY'))
18
  return tokenizer
19
 
20
  def predict(text, model, tokenizer):
@@ -23,7 +23,7 @@ def predict(text, model, tokenizer):
23
  return outputs
24
 
25
  def main():
26
- model_name = os.getenv('MODEL_PATH')
27
  if model_name is None:
28
  raise ValueError("MODEL_PATH environment variable not set or is None")
29
 
 
8
  def load_model(model_name):
9
  try:
10
  # Load TensorFlow model from Hugging Face
11
+ model = TFBertForSequenceClassification.from_pretrained(model_name, use_auth_token=os.getenv('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'), from_tf=True)
12
  except OSError:
13
  raise ValueError("Model loading failed.")
14
  return model
15
 
16
  def load_tokenizer(model_name):
17
+ tokenizer = BertTokenizerFast.from_pretrained(model_name, use_auth_token=os.getenv('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'))
18
  return tokenizer
19
 
20
  def predict(text, model, tokenizer):
 
23
  return outputs
24
 
25
  def main():
26
+ model_name = os.getenv('Erfan11/Neuracraft')
27
  if model_name is None:
28
  raise ValueError("MODEL_PATH environment variable not set or is None")
29