模型效果不错，可是跑的时候遇到显存不释放的问题，推理4-5次就会爆显存，缩短文本长度也只是推迟爆显存的时间。
请问要增加什么代码来释放显存呢？卡是4090，24G
文本长度大概在1000左右

以下是代码
#!/usr/bin/env python

coding: utf-8

In[1]:

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation.utils import GenerationConfig
import ast
import json
from tqdm import tqdm
import gc
torch.backends.cudnn.enabled = True

torch.backends.cudnn.benchmark = True

In[2]:

ckpt = "/root/autodl-tmp/model/Baichuan2-13B-Chat-4bits"
tokenizer = AutoTokenizer.from_pretrained(ckpt,
revision="v2.0",
use_fast=False,
trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(ckpt,
revision="v2.0",
device_map="auto",
torch_dtype=torch.bfloat16,
trust_remote_code=True)
model.generation_config = GenerationConfig.from_pretrained(ckpt, revision="v2.0")
model.eval()
do_infer()

In[3]:

def convert_punctuation(text):
# 中文标点替换成英文标点
punctuation_mapping = {
'，': ',',
'。': '.',
'：': ':',
'；': ';',
'？': '?',
'！': '!',
'【': '[',
'】': ']',
'（': '(',
'）': ')'
}
translation_table = str.maketrans(punctuation_mapping)
return text.translate(translation_table)

In[4]:

#使用大模型进行ner,然后返回字符串,用ast检查格式是否正确,如果不正确修改参数重新尝试,尝试次数三次,三次后依然不正确返回空列表。
#把结果和输入一起写到文本文档里

def check_format(llmner_output, tuple_len = 2):
llmner_output = convert_punctuation(llmner_output)
try:
parsed_output = ast.literal_eval(llmner_output)
if isinstance(parsed_output, list):
for item in parsed_output:
if not isinstance(item, tuple) or len(item) != tuple_len:
print("元组长度错误")
return False
return True
else:
print("文本并非列表")
return False
except (SyntaxError, ValueError):
print("无法转换文本为列表")
return False

test_input1 = "[('乐山电力股份有限公司', '公司类型'), ('乐山乐电天威硅业科技有限责任公司', '公司类型'), ('四川省乐山市中级人民法院', '法院'), ('中国证券报', '报纸'), ('上海证券报', '报纸'), ('证券时报', '报纸'), ('证券日报', '报纸'), ('中国进出口银行', '机构'), ('招银金融租赁有限公司', '机构'), ('中国外贸金融租赁有限公司', '机构')]"
test_input2 = "[('乐山电力股份有限公司', '公司类型'), ('乐山乐电天威硅业科技有限责任公司', '公司类型'), ('四川省乐山市中级人民法院', '法院'), ('中国证券报', '报纸'), ('上海证券报', '报纸'), ('证券时报', '报纸'), ('证券日报', '报纸'), ('中国进出口银行', ), ('招银金融租赁有限公司', '机构'), ('中国外贸金融租赁有限公司', '机构')]"
test_input3 = "[('乐山电力股份有限公司', '公司类型 ('乐山乐电天威硅业科技有限责任公司', '公司类型'), ('四川省乐山市中级人民法院', '法院'), ('中国证券报', '报纸'), ('上海证券报', '报纸'), ('证券时报', '报纸'), ('证券日报', '报纸'), ('中国进出口银行', ), ('招银金融租赁有限公司', '机构'), ('中国外贸金融租赁有限公司', '机构')]"

print(check_format(test_input1))
print(check_format(test_input2))
print(check_format(test_input3))

In[5]:

def process_json_data(file_path):
json_list = []
with open(file_path, 'r') as file:
for line in file:
try:
json_data = json.loads(line)
# 在这里处理JSON数据,例如提取特定字段
# 例如,假设JSON数据具有字段“name”,我们可以打印它
#print(json_data['name'])
json_list.append(json_data)
# 在此处执行其他操作,根据您的需求进行相应的处理

        except json.JSONDecodeError as e:
            print(f"解析JSON数据时出错：{e}")
        except KeyError as e:
            print(f"未找到所需的字段：{e}")
        
return json_list

text_path = '/root/autodl-tmp/ccks2020_subtask2/event_element_train_data_label.txt'
txt_iter = process_json_data(text_path)

In[6]:

txt_iter[0]

In[7]:

def llmner(input_text):

torch.cuda.empty_cache()

model.generation_config = GenerationConfig(
    assistant_token_id= 196,
    bos_token_id= 1,
    do_sample= False,
    eos_token_id= 2,
    min_new_tokens= 5,
    max_new_tokens = 2048,
    pad_token_id= 0,
    num_beams = 4,
    repetition_penalty= 1.2,
    #temperature=0.85,
    #top_k= 20,
    #top_p= 0.85,
    user_token_id= 195,
)

prefix = f"""
从给定的文本中提取出可能的实体和实体类型,以('实体','实体类型')的格式回答
请注意,你的返回格式为一个由('实体','实体类型')对象组成的列表[],其格式为：
\'''
[('实体','实体类型'),('实体','实体类型'),('实体','实体类型'),...]
\'''
如果没有抽取出实体,直接返回[]
"""
#input_text = "证券代码：600644 股票简称：乐山电力 编号：临2016-035乐山电力股份有限公司关于乐电天威公司破产清算进展情况的公告本公司董事会及全体董事保证本公告内容不存在任何虚假记载、误导性陈述或者重大遗漏,并对其内容的真实性、准确性和完整性承担个别及连带责任。2014年12月3日,乐山电力股份有限公司以债权人身份向四川省乐山市中级人民法院(以下简称“乐山中院”)申请对公司持股51%的控股子公司乐山乐电天威硅业科技有限责任公司(以下简称“乐电天威公司”)进行破产清算。乐山中院于2014年12月17日受理乐山电力对乐电天威公司的破产清算申请。2015年3月27日,乐山中院出具《民事裁定书》(2015乐民破字第1-1号),正式宣告乐电天威公司破产。乐电天威公司破产财产于2016年5月6日在乐山产权交易中心第六次整体拍卖,竞买人乐山高新投资建设开发有限公司以18,753.95万元竞得该标的。乐电天威公司破产财产中的渣场资产在第五次竞拍中以800万元成交。(详见公司于2016年3月2日、4月30日和5月21日在上海证券交易所网站和《中国证券报》《上海证券报》《证券时报》《证券日报》上披露的公告)近日,公司收到四川省乐山市中级人民法院《民事裁定书》(2015)乐民破字第1-5号,有关内容如下：乐山中院认为,乐电天威公司管理人召集债权人于2016年5月31日召开第三次债权人会议,审议并通过了管理人制作的《乐山乐电天威硅业科技有限责任公司破产财产分配方案》,乐山中院经审查认为《乐山乐电天威硅业科技有限责任公司破产财产分配方案》符合法律规定,据此,依照《中华人民共和国企业破产法》第一百一十五条之规定,裁定：对第三次债权人会议通过的《乐山乐电天威硅业科技有限责任公司破产财产分配方案》予以认可。该裁定从6月16日起生效。根据《乐山乐电天威硅业科技有限责任公司破产财产分配方案》,公司预计将冲回部分因担保(中国进出口银行、招银金融租赁有限公司、中国外贸金融租赁有限公司)已计提的预计负债约3700万元；同时,公司预计将获得约230万元的债权清偿,并计入当期损益。最终金额以会计师事务所年度审计确认后的结果为准。有关乐电天威公司破产清算事项的进展情况公司将及时履行信息披露义务,敬请广大投资者注意投资风险。备查文件：1.四川省乐山市中级人民法院《民事裁定书》(2015)乐民破字第1-5号；2.《乐山乐电天威硅业科技有限责任公司破产财产分配方案》。特此公告。乐山电力股份有限公司2016年6月21日"

suffix = f"""
文本段落：
{input_text[:512]}
你的回答是：
"""


text = prefix+suffix
messages = []
messages.append({"role": "user", "content": text})
response = model.chat(tokenizer, messages)
return response

In[8]:

def llmtriple(input_text,entities):

torch.cuda.empty_cache()

model.generation_config = GenerationConfig(
    assistant_token_id= 196,
    bos_token_id= 1,
    do_sample= False,
    eos_token_id= 2,
    min_new_tokens= 5,
    max_new_tokens = 2048,
    pad_token_id= 0,
    num_beams = 2,
    repetition_penalty= 1.2,
    #temperature=0.15,
    #top_k= 20,
    #top_p= 0.85,
    user_token_id= 195,
)

prefix = f"""
从给定的文本中提取出给定实体之间可能的关系三元组,以('头实体','关系类别','尾实体')的格式回答
请注意,你的返回格式为一个由('头实体','关系类别','尾实体')关系三元组组成的列表[],其格式为：
\'''
[('头实体','关系类别','尾实体'),('头实体','关系类别','尾实体'),...]
\'''
如果没有抽取出关系,直接返回[]
"""
#input = "证券代码：600644 股票简称：乐山电力 编号：临2016-035乐山电力股份有限公司关于乐电天威公司破产清算进展情况的公告本公司董事会及全体董事保证本公告内容不存在任何虚假记载、误导性陈述或者重大遗漏,并对其内容的真实性、准确性和完整性承担个别及连带责任。2014年12月3日,乐山电力股份有限公司以债权人身份向四川省乐山市中级人民法院(以下简称“乐山中院”)申请对公司持股51%的控股子公司乐山乐电天威硅业科技有限责任公司(以下简称“乐电天威公司”)进行破产清算。乐山中院于2014年12月17日受理乐山电力对乐电天威公司的破产清算申请。2015年3月27日,乐山中院出具《民事裁定书》(2015乐民破字第1-1号),正式宣告乐电天威公司破产。乐电天威公司破产财产于2016年5月6日在乐山产权交易中心第六次整体拍卖,竞买人乐山高新投资建设开发有限公司以18,753.95万元竞得该标的。乐电天威公司破产财产中的渣场资产在第五次竞拍中以800万元成交。(详见公司于2016年3月2日、4月30日和5月21日在上海证券交易所网站和《中国证券报》《上海证券报》《证券时报》《证券日报》上披露的公告)近日,公司收到四川省乐山市中级人民法院《民事裁定书》(2015)乐民破字第1-5号,有关内容如下：乐山中院认为,乐电天威公司管理人召集债权人于2016年5月31日召开第三次债权人会议,审议并通过了管理人制作的《乐山乐电天威硅业科技有限责任公司破产财产分配方案》,乐山中院经审查认为《乐山乐电天威硅业科技有限责任公司破产财产分配方案》符合法律规定,据此,依照《中华人民共和国企业破产法》第一百一十五条之规定,裁定：对第三次债权人会议通过的《乐山乐电天威硅业科技有限责任公司破产财产分配方案》予以认可。该裁定从6月16日起生效。根据《乐山乐电天威硅业科技有限责任公司破产财产分配方案》,公司预计将冲回部分因担保(中国进出口银行、招银金融租赁有限公司、中国外贸金融租赁有限公司)已计提的预计负债约3700万元；同时,公司预计将获得约230万元的债权清偿,并计入当期损益。最终金额以会计师事务所年度审计确认后的结果为准。有关乐电天威公司破产清算事项的进展情况公司将及时履行信息披露义务,敬请广大投资者注意投资风险。备查文件：1.四川省乐山市中级人民法院《民事裁定书》(2015)乐民破字第1-5号；2.《乐山乐电天威硅业科技有限责任公司破产财产分配方案》。特此公告。乐山电力股份有限公司2016年6月21日"
entity_name = [ i[0] for i in entities]

suffix = f"""
文本中的实体包括：
{entity_name}
文本：
{input_text[:512]}
你抽取出的关系三元组是：
"""


text = prefix+suffix
messages = []
messages.append({"role": "user", "content": text})
response = model.chat(tokenizer, messages)
#print(response)
return response

def ner_pipeline(text):

response_ner = llmner(text)
    
if check_format(response_ner, tuple_len=2):
    entities = ast.literal_eval(convert_punctuation(response_ner))
    return entities
else:
    print('ner输出格式错误')
    print(response_ner)
    return []

def triple_pipline(text, entities):

response_triple = llmtriple(text,entities)
    
if check_format(response_triple, tuple_len=3):
    triples = ast.literal_eval(convert_punctuation(response_triple))
    return triples
else:
    print('triple输出格式错误')
    print(response_triple)
    return []

file_path = './results/ner-triple.txt'

with open(file_path, 'a+', encoding = 'utf-8') as file:
file.write('new_exp'+'\n')

with torch.no_grad():
for txt_json in tqdm(txt_iter):
torch.cuda.empty_cache()
gc.collect()
text = txt_json['content']
entities = ner_pipeline(text)

    if len(entities) > 1:
        torch.cuda.empty_cache()
        gc.collect()
        triples = triple_pipline(text, entities)
    else:
        triples = []

    with open(file_path, 'a+', encoding = 'utf-8') as file:
        file.write(str({"doc_id": txt_json['doc_id'], "entity": entities, "triple": triples})+'\n')