network/main.py

import sys
import configparser
import pandas as pd
import torch
from torch import nn
import torch.optim as optim


def main(**param):
    L ={"dashscope":1536, "liandong":1024, "zhipu":1024}
    answerCache = []

    config = configparser.ConfigParser()
    config.read("settings.ini", encoding="utf-8")
    modelName = param['model']
    embeddings = param['embedding']
    mode = param['mode']
    print(modelName, embeddings, mode)

    emName = embeddings + 'embedding'
    embedding_config = dict(config.items(emName))
    ems = __import__('embeddings.%s' % emName,
                     fromlist=['embeddings'])
    string = 'ems.' + emName.capitalize()
    embedding = eval(string)(**embedding_config)

    mds = __import__('models.%s' % modelName,
                     fromlist=['models'])
    string = 'mds.' + modelName.capitalize()
    model = eval(string)(L[embeddings])
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    if mode == 'train':
        train = pd.read_csv('data/train.csv')
        for i in range(20):
            nloss = 0
            for k in range(len(train)):
                va = embedding.getem(train.iloc[k]['question'])
                vb = embedding.getem(train.iloc[k]['answer'])
                if train.iloc[k]['answer'] not in answerCache:
                    answerCache.append(train.iloc[k]['answer'])
                trainTensor = model.prosess(va, vb)
                output = model(trainTensor)
                # 计算损失
                predict = torch.tensor(train.iloc[k]['label']).float()
                predict = predict.reshape([1, 1])
                loss = criterion(output,
                                 predict)
                # 反向传播并更新权重
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                nloss += loss
                if k % 50 == 0:
                    print(i, k, 'done')
            print('one loop done', nloss/len(train))
        torch.save(model, 'models/%s.pth' % modelName)
    if mode == 'test':
        n = 0
        model = torch.load('models/%s.pth' % modelName)
        model.eval()
        test = pd.read_csv('data/test.csv')
        for i in range(len(test)):
            va = embedding.getem(test.iloc[i]['question'])
            vb = embedding.getem(test.iloc[i]['answer'])
            testTensor = model.prosess(va, vb)
            output = model(testTensor)
            if output > 0.5 and test.iloc[i]['label'] == 1:
                n += 1
            if output < 0.5 and test.iloc[i]['label'] == 0:
                n += 1
        print(n/len(test))


if __name__ == '__main__':
    if not len(sys.argv) == 4:
        arg1 = 'cnn'
        arg2 = 'dashscope'
        arg3 = 'train'
    else:
        # 从命令行参数中获取参数值
        arg1 = sys.argv[1]
        arg2 = sys.argv[2]
        arg3 = sys.argv[3]
    main(model=arg1, embedding=arg2, mode=arg3)