Inicio Explorar Ayuda
Iniciar sesión
LLM
/
chat-python
3
0
Fork 0
Archivos Incidencias 0 Pull Requests 0 Wiki
Árbol: a6db0a6aa0
Ramas Etiquetas
chat-python
/
rag
/
load_model.py

load_model.py 1.9 KB
Histórico Raw

123456789101112131415161718192021222324252627282930313233343536 
  1. # from pymilvus import model
    2. 

  2. # import torch
    3. 

  3. # from transformers import AutoTokenizer, AutoModelForSequenceClassification
    4. 

  4. 

  5. # device = "cuda:1" if torch.cuda.is_available() else "cpu"
    6. 

  6. 

  7. # # 使用sentence transformer方式加载模型
    8. 

  8. # embedding_path = r"/work/models/multilingual-e5-large-instruct"  # 线上路径
    9. 

  9. # # embedding_path = r"G:/work/code/models/multilingual-e5-large-instruct/"  # 本地路径
    10. 

  10. # sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(model_name=embedding_path,device=device)
    11. 

  11. 

  12. # # rerank模型
    13. 

  13. # bce_rerank_model_path = r"/work/models/bce-reranker-base_v1"  # 线上路径
    14. 

  14. # # bce_rerank_model_path = r"G:/work/code/models/bce-reranker-base_v1"  # 本地路径
    15. 

  15. # bce_rerank_tokenizer = AutoTokenizer.from_pretrained(bce_rerank_model_path)
    16. 

  16. # bce_rerank_base_model = AutoModelForSequenceClassification.from_pretrained(bce_rerank_model_path).to(device)
    17. 

  17. 

  18. 

  19. from pymilvus import model
    20. 

  20. import torch
    21. 

  21. from transformers import AutoTokenizer, AutoModelForSequenceClassification
    22. 

  22. 

  23. device = "cuda:1" if torch.cuda.is_available() else "cpu"
    24. 

  24. 

  25. # 使用sentence transformer方式加载模型
    26. 

  26. # embedding_path = r"/work/models/multilingual-e5-large-instruct"  # 线上路径
    27. 

  27. embedding_path = "/opt/models/multilingual-e5-large-instruct/"  # jk线上路径
    28. 

  28. # embedding_path = r"G:/work/code/models/multilingual-e5-large-instruct/"  # 本地路径
    29. 

  29. sentence_transformer_ef = model.dense.SentenceTransformerEmbeddingFunction(model_name=embedding_path,device=device)
    30. 

  30. 

  31. # rerank模型
    32. 

  32. # bce_rerank_model_path = r"/work/models/bce-reranker-base_v1"  # 线上路径
    33. 

  33. bce_rerank_model_path = r"/opt/models/bce-reranker-base_v1"    # jk线上路径
    34. 

  34. # bce_rerank_model_path = r"G:/work/code/models/bce-reranker-base_v1"  # 本地路径
    35. 

  35. bce_rerank_tokenizer = AutoTokenizer.from_pretrained(bce_rerank_model_path)
    36. 

  36. bce_rerank_base_model = AutoModelForSequenceClassification.from_pretrained(bce_rerank_model_path).to(device)
    37.