一个用于解释、预测和引导文本模型/数据的Scikit-learn友好库。
同时包含一些用于文本数据入门的实用工具。

可解释建模/引导

模型	参考	输出	描述
Tree-Prompt	🗂️, 🔗, 📄, 📖,	解释 + 引导	生成提示树以引导LLM（官方）
iPrompt	🗂️, 🔗, 📄, 📖	解释 + 引导	生成解释数据模式的提示（官方）
AutoPrompt	ㅤㅤ🗂️, 🔗, 📄	解释 + 引导	使用输入梯度找到自然语言提示
D3	🗂️, 🔗, 📄, 📖	解释	解释两个分布之间的差异
SASC	ㅤㅤ🗂️, 🔗, 📄	解释	使用LLM解释黑盒文本模块（官方）
Aug-Linear	🗂️, 🔗, 📄, 📖	线性模型	使用LLM提取嵌入以拟合更好的线性模型（官方）
Aug-Tree	🗂️, 🔗, 📄, 📖	决策树	使用LLM扩展特征以拟合更好的决策树（官方）
QAEmb	🗂️, 🔗, 📄, 📖	可解释嵌入	通过向LLM提问生成可解释嵌入（官方）
KAN	🗂️, 🔗, 📄, 📖	小型网络	拟合2层Kolmogorov-Arnold网络

📖演示笔记本 🗂️ 文档 🔗 参考代码 📄 研究论文
⌛ 我们计划支持其他可解释算法，如RLPrompt、CBMs和NBDT。如果您想贡献算法，欢迎提交PR 😄

通用工具

模型	参考
🗂️ LLM包装器	轻松调用不同的LLM
🗂️ 数据集包装器	下载经过最小处理的huggingface数据集
🗂️ N元组包	学习N元组的线性模型
🗂️ 线性微调	在LLM嵌入之上微调单个线性层

快速入门

安装: pip install imodelsx（或者，为了更多控制，克隆并从源代码安装）

演示: 请参阅演示笔记本

自然语言解释

Tree-prompt

from imodelsx import TreePromptClassifier
import datasets
import numpy as np
from sklearn.tree import plot_tree
import matplotlib.pyplot as plt

# 设置数据
rng = np.random.default_rng(seed=42)
dset_train = datasets.load_dataset('rotten_tomatoes')['train']
dset_train = dset_train.select(rng.choice(
    len(dset_train), size=100, replace=False))
dset_val = datasets.load_dataset('rotten_tomatoes')['validation']
dset_val = dset_val.select(rng.choice(
    len(dset_val), size=100, replace=False))

# 设置参数
prompts = [
    "这部电影是",
    " 积极还是消极？这部电影是",
    " 这部电影的情感是",
    " 这部电影的剧情真的是",
    " 这部电影的演技是",
]
verbalizer = {0: " 消极。", 1: " 积极。"}
checkpoint = "gpt2"

# 拟合模型
m = TreePromptClassifier(
    checkpoint=checkpoint,
    prompts=prompts,
    verbalizer=verbalizer,
    cache_prompt_features_dir=None,  # 'cache_prompt_features_dir/gp2',
)
m.fit(dset_train["text"], dset_train["label"])


# 计算准确率
preds = m.predict(dset_val['text'])
print('\nTree-Prompt 准确率 (验证集) ->',
      np.mean(preds == dset_val['label']))  # -> 0.7

# 比较单个提示的准确率
for i, prompt in enumerate(prompts):
    print(i, prompt, '->', m.prompt_accs_[i])  # -> 0.65, 0.5, 0.5, 0.56, 0.51

# 可视化决策树
plot_tree(
    m.clf_,
    fontsize=10,
    feature_names=m.feature_names_,
    class_names=list(verbalizer.values()),
    filled=True,
)
plt.show()

iPrompt

from imodelsx import explain_dataset_iprompt, get_add_two_numbers_dataset

# 获取一个简单的两数相加数据集
input_strings, output_strings = get_add_two_numbers_dataset(num_examples=100)
for i in range(5):
    print(repr(input_strings[i]), repr(output_strings[i]))

# 用自然语言提示字符串解释输入和输出之间的关系
prompts, metadata = explain_dataset_iprompt(
    input_strings=input_strings,
    output_strings=output_strings,
    checkpoint='EleutherAI/gpt-j-6B', # 使用哪个语言模型
    num_learned_tokens=3, # 学习多长的提示
    n_shots=3, # 每个示例的shot数
    n_epochs=15, # 搜索多少轮
    verbose=0, # 打印多少信息
    llm_float16=True, # 是否以float16加载模型
)
--------
prompts是找到的自然语言提示字符串列表

D3 (DescribeDistributionalDifferences)

from imodelsx import explain_dataset_d3
hypotheses, hypothesis_scores = explain_dataset_d3(
    pos=positive_samples, # List[str]格式的正面示例
    neg=negative_samples, # 另一个List[str]
    num_steps=100,
    num_folds=2,
    batch_size=64,
)

SASC

这里我们解释一个模块而不是一个数据集

from imodelsx import explain_module_sasc
# 一个对字符串长度做出响应的玩具模块
mod = lambda str_list: np.array([len(s) for s in str_list])

# 一个最长字符串是动物名称的玩具数据集
text_str_list = ["红色", "蓝色", "x", "1", "2", "河马", "大象", "犀牛"]
explanation_dict = explain_module_sasc(
    text_str_list,
    mod,
    ngrams=1,
)

Aug-imodels

像使用scikit-learn模型一样使用这些。在训练过程中，它们通过LLM拟合更好的特征，但在测试时它们非常快速且完全透明。

from imodelsx import AugLinearClassifier, AugTreeClassifier, AugLinearRegressor, AugTreeRegressor
import datasets
import numpy as np

# 设置数据
dset = datasets.load_dataset('rotten_tomatoes')['train']
dset = dset.select(np.random.choice(len(dset), size=300, replace=False))
dset_val = datasets.load_dataset('rotten_tomatoes')['validation']
dset_val = dset_val.select(np.random.choice(len(dset_val), size=300, replace=False))
# 拟合模型
m = AugLinearClassifier(
    checkpoint='textattack/distilbert-base-uncased-rotten-tomatoes',
    ngrams=2, # 使用二元语法
)
m.fit(dset['text'], dset['label'])

# 预测
preds = m.predict(dset_val['text'])
print('验证集准确率', np.mean(preds == dset_val['label']))

# 解释
print('总n元语法系数数量：', len(m.coefs_dict_))
print('最正面的n元语法')
for k, v in sorted(m.coefs_dict_.items(), key=lambda item: item[1], reverse=True)[:8]:
    print('\t', k, round(v, 2))
print('最负面的n元语法')
for k, v in sorted(m.coefs_dict_.items(), key=lambda item: item[1])[:8]:
    print('\t', k, round(v, 2))

KAN

import imodelsx
from sklearn.datasets import make_classification, make_regression
from sklearn.metrics import accuracy_score
import numpy as np

X, y = make_classification(n_samples=5000, n_features=5, n_informative=3)
model = imodelsx.KANClassifier(hidden_layer_size=64, device='cpu',
                               regularize_activation=1.0, regularize_entropy=1.0)
model.fit(X, y)
y_pred = model.predict(X)
print('测试准确率', accuracy_score(y, y_pred))

# 现在尝试回归
X, y = make_regression(n_samples=5000, n_features=5, n_informative=3)
model = imodelsx.kan.KANRegressor(hidden_layer_size=64, device='cpu',
                                  regularize_activation=1.0, regularize_entropy=1.0)
model.fit(X, y)
y_pred = model.predict(X)
print('测试相关性', np.corrcoef(y, y_pred.flatten())[0, 1])

通用工具

简易基准

遵循sklearn API的易于拟合的基准模型。

from imodelsx import LinearFinetuneClassifier, LinearNgramClassifier
# 在LLM嵌入之上拟合一个简单的单层微调模型
m = LinearFinetuneClassifier(
    checkpoint='distilbert-base-uncased',
)
m.fit(dset['text'], dset['label'])
preds = m.predict(dset_val['text'])
acc = (preds == dset_val['label']).mean()
print('验证准确率', acc)

LLM包装器

用于调用不同语言模型的简易API，带有缓存功能（比langchain更轻量）。

import imodelsx.llm
# 支持任何huggingface checkpoint或openai checkpoint（包括对话模型）
llm = imodelsx.llm.get_llm(
    checkpoint="gpt2-xl",  # text-davinci-003, gpt-3.5-turbo, ...
    CACHE_DIR=".cache",
)
out = llm("愿原力与")
llm("愿原力与") # 当再次计算相同字符串时，使用缓存

数据包装器

用于加载huggingface数据集并进行基本预处理的API。

import imodelsx.data
dset, dataset_key_text = imodelsx.data.load_huggingface_dataset('ag_news')
# 确保dset有名为'train'和'validation'的分割，
# 且每个分割的输入数据包含在由{dataset_key_text}给定的列中

imodelsX