Views
No views yet
1# -*- coding: utf-8 -*-
2import safetensors
3from transformers import AutoConfig,AutoTokenizer
4from modeling_bertchunke_zh import BertChunker
5
6# load config and tokenizer
7config = AutoConfig.from_pretrained(
8 "tim1900/bert-chunker-chinese",
9 trust_remote_code=True,
10)
11tokenizer = AutoTokenizer.from_pretrained(
12 "tim1900/bert-chunker-chinese",
13 padding_side="right",
14 model_max_length=config.max_position_embeddings,
15 trust_remote_code=True,
16)
17
18# initialize model
19model = BertChunker(config)
20device='cpu' # or 'cuda'
21model.to(device)
22
23# load tim1900/bert-chunker-chinese/model.safetensors
24state_dict = safetensors.torch.load_file(f"./model.safetensors")
25model.load_state_dict(state_dict)
26
27# text to be chunked
28text='''起点中文网(www.qidian.com)创立于2002年5月,是国内知名的原创文学网站,隶属于阅文集团旗下。起点中文网以推动中国原创文学事业为宗旨,长期致力于原创文学作者的挖掘与培养,并取得了巨大成果:2003年10月,起点中文网开启“在线收费阅读”服务,成为真正意义上的网络文学赢利模式的先锋之一,就此奠定了原创文学的行业基础。此后,起点又推出了作家福利、文学交互、内容发掘推广、版权管理等机制和体系,为原创文学的发展注入了巨大活力,有力推动了中国文学原创事业的发展。在清晨的微光中,一只孤独的猫头鹰在古老的橡树上低声吟唱,它的歌声如同夜色的回声,穿越了时间的迷雾。树叶在微风中轻轻摇曳,仿佛在诉说着古老的故事,每一个音符都带着森林的秘密。一位年轻的程序员正专注地敲打着键盘,代码的海洋在他眼前展开。他的手指在键盘上飞舞,如同钢琴家在演奏一曲复杂的交响乐。屏幕上的光标闪烁,仿佛在等待着下一个指令,引领他进入未知的数字世界。'''
29
30# chunk the text. The lower threshold is, the more chunks will be generated.
31chunks=model.chunk_text(text, tokenizer, threshold=0.5)
32
33# print chunks
34for i, c in enumerate(chunks):
35 print(f'-----chunk: {i}------------')
36 print(c)