import argparse
from fastllm_pytools import llm
import time

def args_parser():
    parser = argparse.ArgumentParser(description = 'fastllm_chat_demo')
    parser.add_argument('-p', '--path', type = str, required = True, default = '', help = '模型文件的路径')
    args = parser.parse_args()
    return args

if __name__ == "__main__":
    args = args_parser()

    model_path = args.path

    prompts = ["深圳有什么好玩的", "上海有什么好玩的", "晚上睡不着怎么办", "南京有什么好吃的"] * 2
    prompts = [
            "23. 已知动点 $P 、 Q$ 都在曲线 $C:\\left\\{\\begin{array}{l}x=2 \\cos \\beta \\\\ y=2 \\sin \\beta\\end{array}\\right.$ ( $\\beta$ 为参数) 上, 对应参数分别为 $\\beta=\\alpha$ 与 $\\beta=2 \\alpha(0<\\alpha<2 \\pi), M$ 为 $P Q$ 的中点.\n\n(1) 求 $M$ 的轨迹的参数方程;\n\n(2) 将 $M$ 到坐标原点的距离 $d$ 表示为 $\\alpha$ 的函数, 并判断 $M$ 的轨迹是否过坐标 原点.\n",
            "23. 在直角坐标系 $x O y$ 中, 曲线 $C_{1}$ 的参数方程为 $\\left\\{\\begin{array}{l}x=2 \\cos \\alpha \\\\ y=2+2 \\sin \\alpha\\end{array}\\right.$ ( $\\alpha$ 为参数) $M$ 是 $\\mathrm{C}_{1}$ 上的动点, $\\mathrm{P}$ 点满足 $\\overrightarrow{\\mathrm{OP}}=2 \\overrightarrow{\\mathrm{OM}}, \\mathrm{P}$ 点的轨迹为曲线 $\\mathrm{C}_{2}$\n\n(I) 求 $\\mathrm{C}_{2}$ 的方程;\n\n(II ) 在以 $O$ 为极点, $x$ 轴的正半轴为极轴的极坐标系中, 射线 $\\theta=\\frac{\\pi}{3}$ 与 $C_{1}$ 的异 于极点的交点为 $A$, 与 $C_{2}$ 的异于极点的交点为 $B$, 求 $|A B|$.\n",
            "阅读下面这首唐诗，完成各题。  \n春日秦国怀古   周朴① \n荒郊一望欲消魂②，泾水萦纡傍远村。  \n牛马放多春草尽，原田耕破古碑存。  \n云和积雪苍山晚，烟伴残阳绿树昏。  \n数里黄沙行客路，不堪回首思秦原。  \n[注]①周朴（？～ 878）：字太朴，吴兴（今属浙江）人。②消魂；这里形容极\n其哀愁。③ 泾水：渭水支流，在今陕西省中部，古属秦国。 萦纡：旋绕曲折 。 \n（1）这首诗表现了诗人什么样的感情？请简要分析。  \n（2）你认为这首诗在写作上是如何处理情景关系的？\n",
            "阅读下面这首宋诗，完成下列各题。  \n次韵雪后书事二首（其一）  \n朱熹  \n  惆怅江头几树梅，杖藜行绕去还来。  \n前时雪压无寻处，昨夜月明依旧开。  \n折寄遥怜人似玉，相思应恨劫成灰。  \n沉吟日落寒鸦起，却望柴荆独自回。  \n（1）这首咏梅诗中，作者用什么手法来表现梅花的？请简要分析  \n \n \n（2）诗的最后一联表达了作者什么样的心情？请简要分析。\n",
            "17. (12 分) 已知 $a, b, c$ 分别为 $\\triangle A B C$ 三个内角 $A, B, C$ 的对边, $c=\\sqrt{3} a \\sin C-\\cos A$.\n\n(1) 求 $A$;\n\n(2) 若 $a=2, \\triangle A B C$ 的面积为 $\\sqrt{3}$, 求 $b, c$.\n",
            "阅读下面这首宋诗，完成下列各题。  \n礼部贡院阅进士就试  欧阳修  \n紫案焚香暖吹轻，广庭清晓席群英。  \n无哗战士衔枚勇，下笔春蚕食叶声。  \n  乡里献贤先德行，朝廷列爵待公卿。  \n自惭衰病心神耗，赖有群公鉴裁精。  \n（1）下列对这首诗的赏析，不恰当的两项是〔 5分〕　   　。 \nA．诗的第一句写出了考场肃穆而又怡人的环境，衬托出作者的喜悦心情。  \nB．第三句重点在表现考生奋勇争先、一往无前，所以把他们比作战士。  \nC．参加礼部考试的考生都由各地选送而来，道德品行是选送的首要依据。  \nD．朝廷对考生寄予了殷切的期望，希望他们能够成长为国家的栋梁之材。  \nE．作者承认自己体弱多病的事实，表示选材工作要依靠其他考官来完成。  \n（2）本诗的第四句 “下笔春蚕食叶声 ”广受后世称道 ，请赏析这一句的精妙之处 。\n〔6分〕\n",
            "阅读下面这首唐诗，完成下列各题。  \n野歌  李贺  \n鸦翎羽箭山桑弓，仰天射落衔芦鸿。  \n麻衣黑肥冲北风，带酒日晚歌田中。  \n男儿屈穷心不穷，枯荣不等嗔天公。  \n寒风又变为春柳，条条看即烟濛濛。  \n（1）下列对这首诗的赏析，不正确的一项是 　   　 \nA．弯弓射鸿、麻衣冲风、饮酒高歌都是诗人排解心头苦闷与抑郁的方式。  \nB．诗人虽不得不接受生活贫穷的命运，但意志并未消沉，气概仍然豪迈。  \nC．诗中形容春柳的方式与韩愈《早春呈水部张十八员外》相同，较为常见。  \nD．本诗前半描写场景，后半感事抒怀，描写与抒情紧密关联，脉络清晰。  \n（2）诗的最后两句有何含意？请简要分析。\n",
            "17. (12 分) 已知 $a, b, c$ 分别是 $\\triangle A B C$ 内角 $A, B, C$ 的对边, $\\sin ^{2} B=2 \\sin A \\sin C$\n\n(I ) 若 $a=b$, 求 $\\cos B$;\n\n（II ) 设 $B=90^{\\circ}$, 且 $a=\\sqrt{2}$, 求 $\\triangle A B C$ 的面积.\n",
            "17. （12 分) 已知等比数列 $\\left\\{a_{n}\\right\\}$ 中, $a_{1}=\\frac{1}{3}$, 公比 $q=\\frac{1}{3}$.\n\n( I ) $S_{n}$ 为 $\\left\\{a_{n}\\right\\}$ 的前 $n$ 项和, 证明: $S_{n}=\\frac{1-a_{n}}{2}$\n\n(II ) 设 $b_{n}=\\log _{3} a_{1}+\\log _{3} a_{2}+\\ldots+\\log _{3} a_{n}$, 求数列 $\\left\\{b_{n}\\right\\}$ 的通项公式.\n",
            "阅读下面这首宋词，完成各题。  \n阮郎归  \n无名氏①  \n春风吹雨绕残枝，落花无可飞。小池寒渌欲生漪，雨晴还日西。  \n帘半卷，燕双归。讳愁无奈眉②．翻身整顿着残棋，沉吟应劫迟③．  \n[注]①作者一作秦观。②讳愁：隐瞒内心的痛苦。③劫：围棋术语。  \n（1）词上半阕的景物描写对全词的感情抒发起了什么作用？请结合内容分析。  \n \n \n（2）末尾两句表现了词人什么样的情绪，是如何表现的，请简要阐述。\n",
            "3．（11 分）古代诗歌阅读 阅读下面的宋诗，完成问题。\n内宴奉诏作 曹翰① 三十年前学六韬②，英名常得预时髦③． 曾因国难披金甲，不为家贫卖宝刀。 臂健尚嫌弓力软，眼明犹识阵云高④． 庭前昨夜秋风起，羞见盘花旧战袍。 【注】①曹翰（923～992），宋初名将，②六韬：古代兵书。 ③时髦： 指当代俊杰。 ④阵云：战争中的云气，这里有站阵之意。 （1）诗的颈联又作“臂弱尚嫌弓力软，眼昏犹识阵云高”，你认为哪一种比较好？ 为什么？请简要分析。 （2）这首诗与辛弃疾的《破阵子（醉里挑灯看剑）》题材相似，但情感基调却有 所不同，请指出二者的不同之处。\n",
            "阅读下面这首乐府诗，完成问题。  \n       雨雪曲       \n江总① \n雨雪隔榆溪②，从军度陇西③．绕阵看狐迹，依山见马蹄。  \n天寒旗彩坏，地暗鼓声低。漫漫愁云起，苍苍别路迷。  \n    【注】①江总（ 518～590）：南朝陈文学家，字总持，济阳考城（今河南兰\n考）人。历仕梁、陈、隋三朝。  ②榆溪：指边塞。③陇西：在今甘肃东部。  \n（1）这首诗描写了什么样的环境？末句中的 “别路 ”是什么意思？  \n \n \n（2）诗人把 “旗彩坏 ”、“鼓声低 ”分别接在 “天寒 ”、“地暗 ”之后，这样写有什么好\n处？这首诗表现了戍卒什么样的情感？\n",
            "阅读下面这首宋词，完成下列各题。  \n鹊 桥 仙 \n陆 游 \n华灯纵博，雕鞍驰射，谁记当年豪举①？酒徒一一取封候，独去作江边渔父。  \n轻舟八尺，低逢三扇，占断苹洲烟雨②．镜湖③元自属闲人，又何必君恩赐与！  \n    【注】①这三句是追忆当年军中的生活。博，古代的一种棋戏。  ②占断：\n占尽。苹洲烟雨：指长满苹草、烟雨空濛的风光。③镜湖：即鉴湖，在今浙\n江绍兴。唐天宝初，贺知章请求回家乡会稽当道士，玄宗诏赐他镜湖一角。  \n（1）上阕最后两句是什么意思？他表达了作者什么意思的情感？  \n \n \n \n（2）词的结尾借用了贺知章的故事，这有什么用意？请简要分析。\n",
            "阅读下面这首宋词，完成下列各题。  \n思远人  \n晏几道红叶黄花秋意晚，千里念行客。飞云过尽，归鸿无信，何处寄书得。            \n泪弹不尽临窗滴，就砚旋研墨。渐写到别来，此情深处，红笺为无色。  \n（1）这首词表达了什么样的感情？ “红叶黄花秋意晚 ”一句对表达这种感情有什\n么作用？  \n \n \n（2）“就砚旋研墨 ”与“临窗滴 ”有什么关系？ “红笺为无色 ”的原因是什么？请简要\n分析。\n",
            "24. 设函数 $f(x)=|x-a|+3 x$, 其中 $a>0$.\n\n(I) 当 $a=1$ 时, 求不等式 $f(x) \\geqslant 3 x+2$ 的解集\n\n(II ) 若不等式 $f(x) \\leqslant 0$ 的解集为 $\\{x \\mid x \\leqslant-1\\}$, 求 $a$ 的值.\n",
            "19. (10 分) 为调查某地区老年人是否需要志愿者提供帮助, 用简单随机抽样 方法从该地区调查了 500 位老年人，结果如表:\n\n\\begin{tabular}{|c|c|c|}\n\\hline 性别 & 男 & 女 \\\\\n\\hline 是否需要志愿者 & & \\\\\n\\hline 需要 & 40 & 30 \\\\\n\\hline 不需要 & 160 & 270 \\\\\n\\hline\n\\end{tabular}\n\n（1）估计该地区老年人中，需要志愿者提供帮助的比例;\n\n(2) 能否有 $99 \\%$ 的把握认为该地区的老年人是否需要志愿者提供帮助与性别有 关?\n\n(3) 根据（2）的结论, 能否提出更好的调查方法来估计该地区的老年人中需要 志愿者提供帮助的老年人比例? 说明理由. \n\n\\begin{tabular}{|c|c|c|c|}\n\\hline$P\\left(K^{2} \\geqslant k\\right)$ & 0.050 & 0.010 & 0.001 \\\\\n\\hline & 3.841 & 6.635 & 10.828 \\\\\n\\hline\n\\end{tabular}\n\n附: $K^{2}=\\frac{n(a d-b c)^{2}}{(a+b)(c+d)(a+c)(b+d)}$.\n"
        ] * 2
    prompts = ["满江红全文"] * 32
    print(prompts)

    responses, historys = [], []
    use_hf_model = False
    if use_hf_model:
        # 以qwen为例
        from transformers import AutoModelForCausalLM, AutoTokenizer
        from transformers.generation import GenerationConfig
        tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cpu", trust_remote_code=True, fp32=True).eval()
        model.generation_config = GenerationConfig.from_pretrained(model_path, trust_remote_code=True) # 可指定不同的生成长度、top_p等相关超参
        llm.set_device_map(["cuda:1", "cuda:2"])
        model = llm.from_hf(model, tokenizer, dtype = "float16")

        t0 = time.time()
        responses, historys = model.chat_batch(tokenizer, prompts)
        t1 = time.time()

    else:
        # llm.set_device_map(["cuda:1", "cuda:2"])
        model = llm.model(model_path)

        t0 = time.time()
        responses, historys = model.response_batch(prompts)
        t1 = time.time()

    token_output_count = 0
    word_len = 0
    for i, res in enumerate(responses):
        tokens = model.tokenizer_encode_string(res)
        token_output_count += len(tokens)
        word_len += len(res)

        print("batch index: ", i)
        print(res)
        print(historys[i])
        print("\n")

    print("\ntokens_count: {}, token/s: {:.2f}, character/s: {:.2f}".format(token_output_count, token_output_count/(t1-t0), word_len/(t1-t0)))