Wisdom Chen commited on
Commit
6740c9f
·
unverified ·
1 Parent(s): 3844cc6

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +20 -8
model.py CHANGED
@@ -51,17 +51,31 @@ def initialize_models() -> bool:
51
  global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
52
 
53
  try:
54
- # CLIP initialization remains the same...
 
 
 
 
 
 
 
 
 
 
 
 
55
 
56
  # Initialize LLM with optimized settings
57
  try:
58
  model_name = "mistralai/Mistral-7B-v0.1"
59
  quantization_config = BitsAndBytesConfig(
60
  load_in_4bit=True,
61
- bnb_4bit_quant_type="nf4",
62
- bnb_4bit_use_double_quant=True
 
63
  )
64
 
 
65
  hf_token = st.secrets.get("HUGGINGFACE_TOKEN")
66
  if not hf_token:
67
  raise ValueError("HUGGINGFACE_TOKEN not found in Streamlit secrets")
@@ -70,19 +84,17 @@ def initialize_models() -> bool:
70
  llm_tokenizer = AutoTokenizer.from_pretrained(
71
  model_name,
72
  token=hf_token,
73
- trust_remote_code=True,
74
- revision="v0.1"
75
  )
76
  llm_tokenizer.pad_token = llm_tokenizer.eos_token
77
-
78
  llm_model = AutoModelForCausalLM.from_pretrained(
79
  model_name,
80
  token=hf_token,
81
  quantization_config=quantization_config,
82
  device_map="auto",
83
  torch_dtype=torch.float16,
84
- trust_remote_code=True,
85
- revision="v0.1"
86
  )
87
  llm_model.eval()
88
  print("LLM initialized successfully")
 
51
  global clip_model, clip_preprocess, clip_tokenizer, llm_tokenizer, llm_model, device
52
 
53
  try:
54
+ print(f"Initializing models on device: {device}")
55
+
56
+ # Initialize CLIP model with error handling
57
+ try:
58
+ clip_model, _, clip_preprocess = open_clip.create_model_and_transforms(
59
+ 'hf-hub:Marqo/marqo-fashionCLIP'
60
+ )
61
+ clip_model = clip_model.to(device)
62
+ clip_model.eval()
63
+ clip_tokenizer = open_clip.get_tokenizer('hf-hub:Marqo/marqo-fashionCLIP')
64
+ print("CLIP model initialized successfully")
65
+ except Exception as e:
66
+ raise RuntimeError(f"Failed to initialize CLIP model: {str(e)}")
67
 
68
  # Initialize LLM with optimized settings
69
  try:
70
  model_name = "mistralai/Mistral-7B-v0.1"
71
  quantization_config = BitsAndBytesConfig(
72
  load_in_4bit=True,
73
+ bnb_4bit_compute_dtype=torch.float16,
74
+ bnb_4bit_use_double_quant=True,
75
+ bnb_4bit_quant_type="nf4"
76
  )
77
 
78
+ # Get token from Streamlit secrets
79
  hf_token = st.secrets.get("HUGGINGFACE_TOKEN")
80
  if not hf_token:
81
  raise ValueError("HUGGINGFACE_TOKEN not found in Streamlit secrets")
 
84
  llm_tokenizer = AutoTokenizer.from_pretrained(
85
  model_name,
86
  token=hf_token,
87
+ trust_remote_code=True
 
88
  )
89
  llm_tokenizer.pad_token = llm_tokenizer.eos_token
90
+
91
  llm_model = AutoModelForCausalLM.from_pretrained(
92
  model_name,
93
  token=hf_token,
94
  quantization_config=quantization_config,
95
  device_map="auto",
96
  torch_dtype=torch.float16,
97
+ trust_remote_code=True
 
98
  )
99
  llm_model.eval()
100
  print("LLM initialized successfully")