Shushmita commited on
Commit
39decae
·
verified ·
1 Parent(s): 3b2f6fb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -5
app.py CHANGED
@@ -5,15 +5,16 @@ import torch
5
  # Load model and tokenizer
6
  @st.cache_resource()
7
  def load_model():
8
- model_name = "deepseek-ai/deepseek-coder-1.3b-instruct" # Smaller version
9
-
10
  tokenizer = AutoTokenizer.from_pretrained(model_name)
 
11
  model = AutoModelForCausalLM.from_pretrained(
12
  model_name,
13
- torch_dtype=torch.float16,
14
- device_map="auto" # Automatically assigns to CPU/GPU if available
15
  )
16
-
17
  return model, tokenizer
18
 
19
  model, tokenizer = load_model()
 
5
  # Load model and tokenizer
6
  @st.cache_resource()
7
  def load_model():
8
+ model_name = "deepseek-ai/deepseek-coder-1.3b-instruct" # Using smaller model
9
+
10
  tokenizer = AutoTokenizer.from_pretrained(model_name)
11
+
12
  model = AutoModelForCausalLM.from_pretrained(
13
  model_name,
14
+ torch_dtype=torch.float32, # Use float32 for CPU (float16 needs GPU)
15
+ device_map={"": "cpu"} # Force model to use CPU
16
  )
17
+
18
  return model, tokenizer
19
 
20
  model, tokenizer = load_model()