修改代码提高模型速度：import re
def count_rows_cols(latex_code):
    print(latex_code)
    try:
        # 查找列数：根据表格行的定义找到表格列标识符，如 |l|c|c|c|c|
        columns_match = re.search(r'\\begin\{tabular\}\{(.*?)\}', latex_code)
        if columns_match:
            columns = columns_match.group(1)
            num_cols = len([c for c in columns if c.isalpha()])
        else:
            num_cols = 0

        # 查找行数：根据 \\ 和 \hline 分隔符统计表格的行数
        rows = re.split(r'\\\\|\\hline', latex_code)
        # 过滤掉空行和注释行
        rows = [row for row in rows if row.strip() and not row.strip().startswith('%')]

        # 进一步处理行分隔符
        num_rows = 0
        for row in rows:
            if '&' in row:
                num_rows += 1

        return num_rows, num_cols
    except Exception as e:
        print(f"An error occurred: {e}")
        return 6, 6  # 返回0,0表示解析失败
import os
import json
from PIL import Image
import pandas as pd

if os.environ.get('DATA_PATH_B'):   # 提交时会选择隐藏的测试数据集路径（A+B榜），数据集的格式与A榜数据相同，但数目不同（5360张）
    base_dir = os.environ.get('DATA_PATH_B')
else:
    base_dir = '/bohr/form-recognition-train-b6y2/v4'   # 示例，把A榜测试数据集路径作为测试集路径，仅开发时挂载A榜数据用于debug
with open(os.path.join(base_dir, 'dataset.json'), 'r') as f:
    data = json.load(f)

with open(os.path.join(base_dir, 'sample_submission.json'), 'r') as f:
    sub = json.load(f)
from transformers import NougatProcessor, VisionEncoderDecoderModel, NougatTokenizerFast,BitsAndBytesConfig
from PIL import Image
import torch
from tqdm import tqdm
from transformers import AutoModelForCausalLM,AutoTokenizer
import torch
torch.backends.cuda.matmul.allow_tf32 = True
processor = NougatProcessor.from_pretrained("/bohr/nougat-small-6yqm/v1/")
modelb = VisionEncoderDecoderModel.from_pretrained("/bohr/nougat-small-6yqm/v1/").to("cuda")
model_path = "/bohr/interLM-dm63/v1/"  # 已下载好并打包为自定义数据集，方便提交时不联网环境也能直接调用


tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True,torch_dtype=torch.float16)
# Set `torch_dtype=torch.float16` to load model in float16, otherwise it will be loaded as float32 and cause OOM Error.
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", torch_dtype=torch.float16,trust_remote_code=True, load_in_4bit=True)
model = model.eval()


device="cuda"
submission = []
print(len(data))
for item in data:
    image_path = os.path.join(base_dir, 'test_images', item["image_path"])
    image = Image.open(image_path)
    pixel_values =processor(image, return_tensors="pt").pixel_values.to("cuda")
    outputs = modelb.generate(
                pixel_values,
                min_length=1,
                max_new_tokens=1000,
                bad_words_ids=[[processor.tokenizer.unk_token_id]],
            )
    sequence= processor.batch_decode(outputs, skip_special_tokens=True,return_dict_in_generate=True,output_scores=True)[0]
    question = [{"role": "user", "content": sequence+f"""This table caption: "{item["caption"]},According to the provided table and its associated textual description, which subject does the content of the table most likely belong to? Choose one from: ['Physics', 'Mathematics', 'Computer Science', 'Quantitative Biology', 'Quantitative Finance', 'Statistics', 'ElectricalEngineeringandSystemsScience', 'Economics']. If you don't know, select only one at random."""}]
    response1, history = model.chat(tokenizer, question, history=[])

    
    import re
    response1 = re.sub(r'(?<!\\)["\']', '', response1)
    print(response1)

    rows, cols = count_rows_cols(sequence)
    
    
    question =[{"role": "user", "content": sequence+f'Question: "{item["question"]}"\noptions[0]: "{item["options"][0]}"\noptions[1]: "{item["options"][1]}"\noptions[2]: "{item["options"][2]}"\noptions[3]: "{item["options"][3]}"\nplease choice the correct answer from options (an int value from 0, 1, 2, 3)'}]
    response3, history = model.chat(tokenizer, question, history=[])
    try:
        
        if response1 not in ('Physics','Mathematics','ComputerScience','QuantitativeBiology','QuantitativeFinance','Statistics','ElectricalEngineeringandSystemsScience','Economics'):
            for cat in ('Physics','Mathematics','ComputerScience','QuantitativeBiology','QuantitativeFinance','Statistics','ElectricalEngineeringandSystemsScience','Economics'):
                if cat in response1:
                    response1 = cat
                    break
            response1 = 'ComputerScience'
        else:
            pass
    except:
        response1 = 'ComputerScience'
    
    try:
        answer = int(response3[0])
    except:
        answer = 0

    sub_item = {
        "image_path": item["image_path"],
        "category": str(response1),
        "cols": cols,
        "rows": rows,
        "answer": answer,
    }
    submission.append(sub_item)
    print(sub_item)

with open('submission.json', 'w') as f:
    json.dump(submission, f)