Spaces:
Runtime error
Runtime error
Upload app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,7 @@ import gradio as gr
|
|
3 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
4 |
import torch
|
5 |
|
6 |
-
model_name = "
|
7 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
8 |
|
9 |
chip_map= {
|
@@ -12,8 +12,8 @@ chip_map= {
|
|
12 |
'gpt_neox.final_layer_norm': 0,
|
13 |
'embed_out': 0
|
14 |
}
|
15 |
-
|
16 |
-
model = AutoModelForCausalLM.from_pretrained(model_name)
|
17 |
|
18 |
|
19 |
def predict(input, history=[], MAX_NEW_TOKENS = 500):
|
@@ -49,4 +49,4 @@ with gr.Blocks() as demo:
|
|
49 |
|
50 |
if __name__ == "__main__":
|
51 |
# demo.launch(debug=True, server_name="0.0.0.0", server_port=9991)
|
52 |
-
demo.launch()
|
|
|
3 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
4 |
import torch
|
5 |
|
6 |
+
model_name = "zirui3/gpt_1.4B_oa_instruct"
|
7 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
8 |
|
9 |
chip_map= {
|
|
|
12 |
'gpt_neox.final_layer_norm': 0,
|
13 |
'embed_out': 0
|
14 |
}
|
15 |
+
model = AutoModelForCausalLM.from_pretrained(name, device_map=chip_map, torch_dtype=torch.float16, load_in_8bit=True)
|
16 |
+
#model = AutoModelForCausalLM.from_pretrained(model_name)
|
17 |
|
18 |
|
19 |
def predict(input, history=[], MAX_NEW_TOKENS = 500):
|
|
|
49 |
|
50 |
if __name__ == "__main__":
|
51 |
# demo.launch(debug=True, server_name="0.0.0.0", server_port=9991)
|
52 |
+
demo.launch()
|