python中,如何实现加载一次模型,后续在不同py文件中调用模型,这个模型本身是用sentencetransformer加载的
时间: 2024-12-21 21:22:12 浏览: 4
在 Python 中,如果你想要在一个模块(通常是单独的一个 .py 文件)中加载 Sentencetransformer 模型,并允许其他 .py 文件通过导入共享该模型,你可以采用以下策略:
1. **将模型保存到硬盘**:
首先,在主模块(比如 model_loader.py)里训练好并保存模型。使用 `transformers` 的 `save_pretrained` 函数,例如:
```python
from sentence_transformers import SentenceTransformer
# 训练好的 SentenceTransformer 实例
model = SentenceTransformer('path/to/model')
# 保存模型
model.save_pretrained('./saved_model')
```
2. **创建一个导入模块**:
创建一个名为 `model.py` 的文件,里面只包含模型加载部分:
```python
# model.py
import os
from transformers import SentenceTransformer
# 加载已保存的模型
class SharedSentenceTransformer:
def __init__(self):
self.model_path = './saved_model'
self.model = SentenceTransformer(self.model_path)
shared_model = SharedSentenceTransformer()
```
3. **在其他 .py 文件中导入并使用模型**:
其他需要使用的文件(如 util.py 或 main.py)只需导入 `model.py` 中的 `SharedSentenceTransformer` 类即可:
```python
from model import shared_model
# 使用加载的模型
embeddings = shared_model.model.encode(['This is a test sentence.', 'Another one'])
```
这样,你只需要在 `model_loader.py` 中加载和保存一次模型,然后所有引用 `model.py` 的地方都可以直接使用共享的 `SharedSentenceTransformer` 对象了。这样做的好处是可以节省内存,并确保对模型的修改在所有使用点都会同步。
阅读全文