cutechicken commited on
Commit
4c60e0e
·
verified ·
1 Parent(s): f13b3f1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -3
app.py CHANGED
@@ -19,7 +19,7 @@ class ModelManager:
19
  self.model = None
20
  self.tokenizer = None
21
  self.setup_model()
22
-
23
  def setup_model(self):
24
  try:
25
  self.tokenizer = AutoTokenizer.from_pretrained(
@@ -32,10 +32,23 @@ class ModelManager:
32
  token=HF_TOKEN,
33
  torch_dtype=torch.float16,
34
  device_map="auto",
35
- trust_remote_code=True
 
36
  )
37
  except Exception as e:
38
- raise Exception(f"Model loading failed: {e}")
 
 
 
 
 
 
 
 
 
 
 
 
39
 
40
  class ChatHistory:
41
  def __init__(self):
 
19
  self.model = None
20
  self.tokenizer = None
21
  self.setup_model()
22
+
23
  def setup_model(self):
24
  try:
25
  self.tokenizer = AutoTokenizer.from_pretrained(
 
32
  token=HF_TOKEN,
33
  torch_dtype=torch.float16,
34
  device_map="auto",
35
+ trust_remote_code=True,
36
+ low_cpu_mem_usage=True
37
  )
38
  except Exception as e:
39
+ print(f"Error loading model: {e}")
40
+ # Fallback to basic loading without device_map
41
+ try:
42
+ self.model = AutoModelForCausalLM.from_pretrained(
43
+ MODEL_ID,
44
+ token=HF_TOKEN,
45
+ torch_dtype=torch.float16,
46
+ trust_remote_code=True
47
+ )
48
+ except Exception as e:
49
+ raise Exception(f"Model loading failed completely: {e}")
50
+
51
+
52
 
53
  class ChatHistory:
54
  def __init__(self):