()
| 45 | |
| 46 | @st.cache_resource |
| 47 | def load_model(): |
| 48 | config = MossConfig.from_pretrained(args.model_name) |
| 49 | tokenizer = MossTokenizer.from_pretrained(args.model_name) |
| 50 | if num_gpus > 1: |
| 51 | model_path = args.model_name |
| 52 | if not os.path.exists(args.model_name): |
| 53 | model_path = snapshot_download(args.model_name) |
| 54 | print("Waiting for all devices to be ready, it may take a few minutes...") |
| 55 | with init_empty_weights(): |
| 56 | raw_model = MossForCausalLM._from_config(config, torch_dtype=torch.float16) |
| 57 | raw_model.tie_weights() |
| 58 | model = load_checkpoint_and_dispatch( |
| 59 | raw_model, model_path, device_map="auto", no_split_module_classes=["MossBlock"], dtype=torch.float16 |
| 60 | ) |
| 61 | else: # on a single gpu |
| 62 | model = MossForCausalLM.from_pretrained(args.model_name).half().cuda() |
| 63 | |
| 64 | return tokenizer, model |
| 65 | |
| 66 | |
| 67 | if "history" not in st.session_state: |
no outgoing calls
no test coverage detected