rasyosef commited on
Commit
d982401
1 Parent(s): 2c96457

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -9,7 +9,7 @@ checkpoint = "microsoft/phi-2"
9
 
10
  # Download and load model and tokenizer
11
  tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
12
- model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.float32, device_map="auto", trust_remote_code=True)
13
 
14
  # Streamer
15
  streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_prompt=True)
@@ -22,7 +22,7 @@ phi2 = pipeline(
22
  streamer=streamer,
23
  pad_token_id=tokenizer.eos_token_id,
24
  eos_token_id=tokenizer.eos_token_id,
25
- device_map="auto"
26
  )
27
 
28
  # Function that accepts a prompt and generates text using the phi2 pipeline
 
9
 
10
  # Download and load model and tokenizer
11
  tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
12
+ model = AutoModelForCausalLM.from_pretrained(checkpoint, torch_dtype=torch.float32, device_map="cpu", trust_remote_code=True)
13
 
14
  # Streamer
15
  streamer = TextIteratorStreamer(tokenizer=tokenizer, skip_prompt=True)
 
22
  streamer=streamer,
23
  pad_token_id=tokenizer.eos_token_id,
24
  eos_token_id=tokenizer.eos_token_id,
25
+ device_map="cpu"
26
  )
27
 
28
  # Function that accepts a prompt and generates text using the phi2 pipeline