import requests
import json

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 writeData(dataList,fileName):
    with open(fileName, 'w', encoding='utf-8') as f:
        for data in dataList:
            f.write(data + "\n")



def geneData(promptModel,categoryList,countPerCate):
    geneDataList = []
    for cate in categoryList:
        print(cate)
        prompt = promptModel.format(countPerCate,cate)

        success = False
        while not success:
            success, msg, result = callServer(prompt)
            print(result)
            if success :
                dataList = json.loads(result)
                if len(dataList) > countPerCate-3:
                    for data in dataList:
                        geneDataList.append(cate+"\t"+data.strip()+"\t")
            else:
                success = False
    writeData(geneDataList, "geneData.txt")
    import random
    random.shuffle(geneDataList)
    writeData(geneDataList, "geneDataRandom.txt")

if __name__ == "__main__":
    categoryList = ["电脑","水果","平板","衣服","酒店","洗浴","书籍","蒙牛","电器","手机"]
    promptModel = """你是一位电商评论文本的分类专家，需要生成{}条{}分类的语料，尽可能覆盖该分类的各个方面，每条语料长度不小于30字，不要提及产品的具体名称，开头不要写数字序号，语料不能用于【分类体系】中的其他分类。只允许按照【指定格式】输出，不允许生成任何其他字符。
                【分类体系】
                ["电脑","水果","平板","衣服","酒店","洗浴","书籍","蒙牛","电器","手机"]
                【指定格式】
                ["评论文本1","评论文本2","评论文本3"]"""

    geneData(promptModel, categoryList, 40)