vsrinivas commited on
Commit
2c98ceb
·
verified ·
1 Parent(s): 9f9775c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -3,12 +3,14 @@ import transformers
3
  import torch
4
  import gradio as gr
5
 
 
 
6
  # checkpoint = "tiiuae/falcon-40b-instruct"
7
- # checkpoint ="tiiuae/falcon-7b-instruct"
8
  # checkpoint = "tiiuae/falcon-7b"
9
  # checkpoint = "mistralai/Mixtral-8x7B-Instruct-v0.1"
10
  # checkpoint = "gpt2"
11
- checkpoint = "amazon/FalconLite2"
12
 
13
  model = AutoModelForCausalLM.from_pretrained(
14
  checkpoint, device_map="auto",
 
3
  import torch
4
  import gradio as gr
5
 
6
+ desired_dtype = torch.bfloat16
7
+
8
  # checkpoint = "tiiuae/falcon-40b-instruct"
9
+ checkpoint ="tiiuae/falcon-7b-instruct"
10
  # checkpoint = "tiiuae/falcon-7b"
11
  # checkpoint = "mistralai/Mixtral-8x7B-Instruct-v0.1"
12
  # checkpoint = "gpt2"
13
+ # checkpoint = "amazon/FalconLite2"
14
 
15
  model = AutoModelForCausalLM.from_pretrained(
16
  checkpoint, device_map="auto",