import json
import requests

def callServer(prompt, model="deepseek-ai/DeepSeek-V3.2"):
    url = "https://api.siliconflow.cn/v1/chat/completions"

    headers = {
        "Authorization": f"Bearer sk-envatqyeffxlrtalgexbrueftyzjcezorhhenlexsioqesox",
        "Content-Type": "application/json"
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        #注意，这里的tempreture要调低，保证结果稳定性
        "temperature": 0.1,
        "max_tokens": 6000
    }

    success = True
    msg = "OK"
    think = ""
    result = ""
    try:
        response = requests.post(url, json=payload, headers=headers)
        response.raise_for_status()
        result = response.json()['choices'][0]['message']['content']
    except Exception as e:
        success = False
        msg = "API调用异常:" + str(e)
    return success, msg, result

def loadRawDataSet(fileName):
    dataList = []
    with open(fileName, "r", encoding="utf-8") as file:
        while True:
            line = file.readline()
            if not line:  # 读取到文件末尾时退出
                break
            line = line.strip()
            if len(line)>0:
                dataList.append(line.split("\t",1))
    return dataList

def writeData(dataList,fileName):
    with open(fileName, 'w', encoding='utf-8') as f:
        for data in dataList:
            f.write(data + "\n")

def tagCorrect(dataList,promptModel):
    batchSize = 10
    batchCount = len(dataList)//batchSize

    correctList = []

    for i in range(batchCount+1):
        print(i*batchSize)
        tempStr = ""
        tempList = dataList[i*batchSize:(i+1)*batchSize]
        for data in tempList:
            tempData = {"分类":data[0],"文本":data[1]}
            tempStr +=  json.dumps(tempData,ensure_ascii=False)+"\n"
        tempStr = tempStr.strip()
        if len(tempStr)>0:
            prompt = promptModel+tempStr
            success = False
            while not success:
                success, msg, result = callServer(prompt)
                print(result)
                if success and len(result) == batchSize:
                    correctList += list(result)
                else:
                    success = False
    writeData(correctList,"correct.txt")


if __name__ == "__main__":
    dataList = loadRawDataSet("total.txt")

    promptModel = """你是一位电商评论文本的分类专家，判断在【分类体系】规定的类别中，【输入文本】中的文本对应分类是否正确，并按照【指定格式】输出。输入文本的每一行是一条数据，需要单独判定。
            【分类体系】
            ["电脑","水果","平板","衣服","酒店","洗浴","书籍","蒙牛","电器","手机"]
            【指定格式】
            是否是是是否否是是否
            【输入文本】
            """

    tagCorrect(dataList,promptModel)