sitammeur commited on
Commit
e96ebfa
·
verified ·
1 Parent(s): 9e0886b

Update src/app/model.py

Browse files
Files changed (1) hide show
  1. src/app/model.py +5 -2
src/app/model.py CHANGED
@@ -43,13 +43,16 @@ def load_model_and_tokenizer(model_name: str, device: str) -> Any:
43
  tokenizer = AutoTokenizer.from_pretrained(
44
  model_name, trust_remote_code=True, token=access_token
45
  )
 
 
 
46
  model.eval()
47
 
48
  # Log the successful loading of the model and tokenizer
49
  logging.info("Model and tokenizer loaded successfully.")
50
 
51
- # Return the model and tokenizer
52
- return model, tokenizer
53
 
54
  # Handle exceptions that may occur during model and tokenizer loading
55
  except Exception as e:
 
43
  tokenizer = AutoTokenizer.from_pretrained(
44
  model_name, trust_remote_code=True, token=access_token
45
  )
46
+ processor = AutoProcessor.from_pretrained(
47
+ model_name, trust_remote_code=True, token=access_token
48
+ )
49
  model.eval()
50
 
51
  # Log the successful loading of the model and tokenizer
52
  logging.info("Model and tokenizer loaded successfully.")
53
 
54
+ # Return the model, tokenizer and processor
55
+ return model, tokenizer, processor
56
 
57
  # Handle exceptions that may occur during model and tokenizer loading
58
  except Exception as e: