Wisdom Chen commited on
Commit
93a6683
·
unverified ·
1 Parent(s): d816efe

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +12 -12
model.py CHANGED
@@ -47,12 +47,6 @@ text_faiss: Optional[object] = None
47
  image_faiss: Optional[object] = None
48
 
49
  def initialize_models() -> bool:
50
- """
51
- Initialize CLIP and LLM models with proper error handling and GPU optimization.
52
-
53
- Returns:
54
- bool: True if initialization successful, raises RuntimeError otherwise
55
- """
56
  global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
57
 
58
  try:
@@ -72,21 +66,26 @@ def initialize_models() -> bool:
72
 
73
  # Initialize LLM with optimized settings
74
  try:
75
- if "HF_TOKEN" in os.environ:
76
- login(token=os.environ["HF_TOKEN"])
 
 
 
 
 
77
  model_name = "mistralai/Mistral-7B-v0.1"
78
  quantization_config = BitsAndBytesConfig(
79
  load_in_4bit=True,
80
  bnb_4bit_compute_dtype=torch.float16,
81
  bnb_4bit_use_double_quant=True,
82
- bnb_4bit_quant_type="nf4",
83
- token=os.environ.get("HF_TOKEN")
84
  )
85
 
86
  llm_tokenizer = AutoTokenizer.from_pretrained(
87
  model_name,
88
  padding_side="left",
89
- truncation_side="left"
 
90
  )
91
  llm_tokenizer.pad_token = llm_tokenizer.eos_token
92
 
@@ -94,7 +93,8 @@ def initialize_models() -> bool:
94
  model_name,
95
  quantization_config=quantization_config,
96
  device_map="auto",
97
- torch_dtype=torch.float16
 
98
  )
99
  llm_model.eval()
100
  print("LLM initialized successfully")
 
47
  image_faiss: Optional[object] = None
48
 
49
  def initialize_models() -> bool:
 
 
 
 
 
 
50
  global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
51
 
52
  try:
 
66
 
67
  # Initialize LLM with optimized settings
68
  try:
69
+ # Check for HF_TOKEN and authenticate
70
+ hf_token = os.environ.get("HF_TOKEN")
71
+ if not hf_token:
72
+ raise RuntimeError("HF_TOKEN environment variable is not set")
73
+
74
+ login(token=hf_token)
75
+
76
  model_name = "mistralai/Mistral-7B-v0.1"
77
  quantization_config = BitsAndBytesConfig(
78
  load_in_4bit=True,
79
  bnb_4bit_compute_dtype=torch.float16,
80
  bnb_4bit_use_double_quant=True,
81
+ bnb_4bit_quant_type="nf4"
 
82
  )
83
 
84
  llm_tokenizer = AutoTokenizer.from_pretrained(
85
  model_name,
86
  padding_side="left",
87
+ truncation_side="left",
88
+ token=hf_token
89
  )
90
  llm_tokenizer.pad_token = llm_tokenizer.eos_token
91
 
 
93
  model_name,
94
  quantization_config=quantization_config,
95
  device_map="auto",
96
+ torch_dtype=torch.float16,
97
+ token=hf_token
98
  )
99
  llm_model.eval()
100
  print("LLM initialized successfully")