model_config.py 773 B

1234567891011121314151617181920212223242526272829
  1. import torch.cuda
  2. import torch.backends
  3. embedding_model_dict = {
  4. "ernie-tiny": "nghuyong/ernie-3.0-nano-zh",
  5. "ernie-base": "nghuyong/ernie-3.0-base-zh",
  6. "text2vec": "GanymedeNil/text2vec-large-chinese",
  7. }
  8. # Embedding model name
  9. EMBEDDING_MODEL = "text2vec"
  10. # Embedding running device
  11. EMBEDDING_DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
  12. # supported LLM models
  13. llm_model_dict = {
  14. "chatglm-6b-int4-qe": "THUDM/chatglm-6b-int4-qe",
  15. "chatglm-6b-int4": "THUDM/chatglm-6b-int4",
  16. "chatglm-6b": "THUDM/chatglm-6b",
  17. }
  18. # LLM model name
  19. LLM_MODEL = "chatglm-6b"
  20. # LLM running device
  21. LLM_DEVICE = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"