import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer, BertModel, BertTokenizer
from llmkt import *

# 处理多组文本输入的函数
def process_text_groups(llm4kt, text_groups):
    outputs = []
    for texts in text_groups:
        output = llm4kt(texts)
        outputs.append(output)
    return outputs

# 测试加载
if __name__ == "__main__":
    # 加载Bert模型和分词器
    tokenizer, model = load_bert_model()
    
    # 初始化LLM4KT模型
    llm4kt = LLM4KT(item_model=model, item_tokenizer=tokenizer, user_model=model)
    
    # 第一组测试文本（10句话）
    texts_10 = [
        "Hello, world![ITEM]",
        "How are you?[ITEM]",
        "This is a test.[ITEM]",
        "Let's see how it works.[ITEM]",
        "Transformers are powerful.[ITEM]",
        "Natural Language Processing is fun.[ITEM]",
        "Machine learning is fascinating.[ITEM]",
        "Deep learning models are complex.[ITEM]",
        "Bert is a popular model.[ITEM]",
        "This is the last sentence.[ITEM]"
    ]
    
    # 第二组测试文本（20句话）
    texts_20 = texts_10 + [
        "Another sentence to test.[ITEM]",
        "More data for the model.[ITEM]",
        "Testing with more sentences.[ITEM]",
        "Increasing the input size.[ITEM]",
        "Handling larger inputs.[ITEM]",
        "Scaling up the model.[ITEM]",
        "Evaluating performance.[ITEM]",
        "Checking for consistency.[ITEM]",
        "Ensuring robustness.[ITEM]",
        "Final sentence for testing.[ITEM]"
    ]
    
    # 多组文本输入
    text_groups = [texts_10, texts_20]
    
    # 获取LLM4KT模型的输出
    outputs = process_text_groups(llm4kt, text_groups)
    
    # 打印输出
    for i, output in enumerate(outputs):
        print(f"Outputs for text group {i + 1}:")
        print(output)


