search

wangyuxinwhy / uniem

unified embedding model
Apache License 2.0
826 stars 64 forks source link

MINOR: Support explicit specification of device #87

Closed ncliang closed 1 year ago

ncliang commented 1 year ago

The PR allows for specifying one of cpu, cuda, or mps to be the device type. It allows for quick and easy benchmarking on Mac devices running an M series chips. If not specified, the device chosen defaults to what would have been chosen in SentenceTransformer or similar modules. To benchmark on an M1 mac, one can do something such as

python run_mteb_zh.py  --model-type sentence_transformer --model-id ~/HamaModel --device=mps

Since typer is used, the help text is also updated accordingly

$ python run_mteb_zh.py --help
                                                                                                                         Usage: run_mteb_zh.py [OPTIONS]                                                                                                                                                                                                                _─ Options ────────────────────────────────────────────────────────────────────────────────────────────────────────────_
│ *  --model-type           [sentence_transformer|text2vec|luotuo|erlang  [default: None] [required]                   │
│                           shen|openai|minimax|azure]                                                                 │
│    --model-id             TEXT                                          [default: None]                              │
│    --task-type            [Classification|Reranking|Retrieval|PairClas  [default: TaskType.Classification]           │
│                           sification|All]                                                                            │
│    --task-name            TEXT                                          [default: None]                              │
│    --output-folder        PATH                                          [default: results]                           │
│    --device               [cpu|cuda|mps]                                [default: None]                              │
│    --help                                                               Show this message and exit.                  │
_──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────_
wangyuxinwhy commented 1 year ago

Thank you very much for your contribution