1 Star 0 Fork 15

珠海_JadeLuo/MOSS

forked from Gitee 极速下载/MOSS 
加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
utils.py 485 Bytes
一键复制 编辑 原始数据 按行查看 历史
Tianxiang Sun 提交于 2023-04-25 15:22 . Create utils.py
import torch
from transformers import StoppingCriteria
class StopWordsCriteria(StoppingCriteria):
def __init__(self, stop_indices: list):
self.stop_indices = stop_indices
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
# do not support batch inference
for i in range(len(self.stop_indices)):
if self.stop_indices[-1-i] != input_ids[0][-1-i]:
return False
return True
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/jadeluo/MOSS.git
git@gitee.com:jadeluo/MOSS.git
jadeluo
MOSS
MOSS
main

搜索帮助