import json
from dataclasses import dataclass
from enum import Enum
from typing import List, Dict, Tuple, Literal

class Roles(Enum):
    system = "system"
    user = "user"
    assistant = "assistant"
    tool = "tool"

class MessagesFormatterType(Enum):
    """
    Enum representing different types of predefined messages formatters.
    """

    MISTRAL = 1

@dataclass
class PromptMarkers:
    start: str
    end: str

class MessagesFormatter:
    def __init__(
            self,
            pre_prompt: str,
            prompt_markers: Dict[Roles, PromptMarkers],
            include_sys_prompt_in_first_user_message: bool,
            default_stop_sequences: List[str],
            use_user_role_for_function_call_result: bool = True,
            strip_prompt: bool = True,
            bos_token: str = "<s>",
            eos_token: str = "</s>"
    ):
        self.pre_prompt = pre_prompt
        self.prompt_markers = prompt_markers
        self.include_sys_prompt_in_first_user_message = include_sys_prompt_in_first_user_message
        self.default_stop_sequences = default_stop_sequences
        self.use_user_role_for_function_call_result = use_user_role_for_function_call_result
        self.strip_prompt = strip_prompt
        self.bos_token = bos_token
        self.eos_token = eos_token
        self.added_system_prompt = False

    def get_bos_token(self) -> str:
        return self.bos_token

    def format_conversation(
            self,
            messages: List[Dict[str, str]],
            response_role: Literal[Roles.user, Roles.assistant] | None = None,
    ) -> Tuple[str, Roles]:
        formatted_messages = self.pre_prompt
        last_role = Roles.assistant
        self.added_system_prompt = False
        for message in messages:
            role = Roles(message["role"])
            content = self._format_message_content(message["content"], role)

            if role == Roles.system:
                formatted_messages += self._format_system_message(content)
                last_role = Roles.system
            elif role == Roles.user:
                formatted_messages += self._format_user_message(content)
                last_role = Roles.user
            elif role == Roles.assistant:
                formatted_messages += self._format_assistant_message(content)
                last_role = Roles.assistant
            elif role == Roles.tool:
                formatted_messages += self._format_tool_message(content)
                last_role = Roles.tool

        return self._format_response(formatted_messages, last_role, response_role)

    def _format_message_content(self, content: str, role: Roles) -> str:
        if self.strip_prompt:
            return content.strip()
        return content

    def _format_system_message(self, content: str) -> str:
        formatted_message = self.prompt_markers[Roles.system].start + content + self.prompt_markers[Roles.system].end
        self.added_system_prompt = True
        if self.include_sys_prompt_in_first_user_message:
            formatted_message = self.prompt_markers[Roles.user].start + formatted_message
        return formatted_message

    def _format_user_message(self, content: str) -> str:
        if self.include_sys_prompt_in_first_user_message and self.added_system_prompt:
            self.added_system_prompt = False
            return content + self.prompt_markers[Roles.user].end
        return self.prompt_markers[Roles.user].start + content + self.prompt_markers[Roles.user].end

    def _format_assistant_message(self, content: str) -> str:
        return self.prompt_markers[Roles.assistant].start + content + self.prompt_markers[Roles.assistant].end

    def _format_tool_message(self, content: str) -> str:
        if isinstance(content, list):
            content = "\n".join(json.dumps(m, indent=2) for m in content)
        if self.use_user_role_for_function_call_result:
            return self._format_user_message(content)
        else:
            return self.prompt_markers[Roles.tool].start + content + self.prompt_markers[Roles.tool].end

    def _format_response(
            self,
            formatted_messages: str,
            last_role: Roles,
            response_role: Literal[Roles.user, Roles.assistant] | None = None,
    ) -> Tuple[str, Roles]:
        if response_role is None:
            response_role = Roles.assistant if last_role != Roles.assistant else Roles.user

        prompt_start = self.prompt_markers[response_role].start.strip() if self.strip_prompt else self.prompt_markers[
            response_role].start
        return formatted_messages + prompt_start, response_role

mixtral_prompt_markers = {
    Roles.system: PromptMarkers("", """\n\n"""),
    Roles.user: PromptMarkers("""[INST] """, """ [/INST]"""),
    Roles.assistant: PromptMarkers("""""", """</s>"""),
    Roles.tool: PromptMarkers("", ""),
}

mixtral_formatter = MessagesFormatter(
    "",
    mixtral_prompt_markers,
    True,
    ["</s>"],
)

from transformers import TextStreamer, AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3",)
mis_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", load_in_4bit = True)
mis_model = PeftModel.from_pretrained(mis_model,
                                "svjack/Genshin_Impact_Mistral_v3_Plot_Chat_roleplay_chat_lora_small"
                                )
mis_model = mis_model.eval()

streamer = TextStreamer(tokenizer)

def mistral_hf_predict(messages, mis_model = mis_model,
    tokenizer = tokenizer, streamer = streamer,
    do_sample = True,
    top_p = 0.95,
    top_k = 40,
    max_new_tokens = 512,
    max_input_length = 3500,
    temperature = 0.9,
    repetition_penalty = 1.0,
    device = "cuda"):

    #encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
    #model_inputs = encodeds.to(device)
    prompt, _ = mixtral_formatter.format_conversation(messages)
    model_inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)

    generated_ids = mis_model.generate(model_inputs, max_new_tokens=max_new_tokens,
                                do_sample=do_sample,
                                  streamer = streamer,
                                  top_p = top_p,
                                  top_k = top_k,
                                  temperature = temperature,
                                  repetition_penalty = repetition_penalty,
                                  )
    out = tokenizer.batch_decode(generated_ids)[0].split("[/INST]")[-1].replace("</s>", "").strip()
    return out

out = mistral_hf_predict([
            {
                "role": "system",
                "content": '''
                故事背景:图书管理员丽莎与助手派蒙在寻找偷书者的冒险中交流，揭示了真相并处理了书籍问题。
                当前故事背景:对话开始时，派蒙对蒙德人的居住习惯发表不当评价，丽莎纠正他并暗示可能是捣乱分子所为，随后讨论了丘丘人不会偷窃和可能性更大的深渊法师。在解开封印后，他们进入遗迹，并决定继续深入调查。
                参与者1:丽莎
                参与者1角色经历:丽莎，作为蒙德城南风之狮庙宇的图书管理员，以其严肃认真的工作态度和对书籍的热爱，与旅行者派蒙共同解决图书丢失的问题。她运用元素感知力帮助找寻线索，与伙伴们互动，展现智慧和勇气，同时对偷书者的行为有着坚定的立场，通过惩罚计划来维护图书的尊严。在游戏中，她不仅提供历史背景，还作为知识库，帮助旅行者理解元素和蒙德的历史，她的存在对解决故事中的谜题和对抗敌人至关重要。在蒙德骑士团中，丽莎也协助凯亚和琴，展现她的团队精神和对守护者的责任感。
                参与者1性格特征:丽莎性格严谨，热爱工作，尊重他人，对待偷书者的行为表现出坚定和公正。她聪明且勇敢，善于使用元素感知力解决问题，同时具有深厚的历史知识和对‘四风守护’的理解。她的智慧和责任感在剧情中起到了关键作用。
                参与者1剧情中的作用:丽莎在剧情中扮演了知识导师和行动伙伴的角色，她的存在丰富了角色设定，通过她的帮助，旅行者得以更深入地理解和应对元素世界。她的出现推动了故事的发展，通过她的智慧和勇气，解决了许多难题，强化了角色间的互动和团队合作。同时，她的责任感和对蒙德的热爱也深化了游戏的主题，体现了对守护者的尊重和对家乡的忠诚。
                参与者2:派蒙
                参与者2角色经历:派蒙是宵宫的旅伴，他们共同面对船的问题，逃离稻妻，与森彦、阿创等人互动，参与烟花制作，解决孩子们的误解。派蒙对宵宫的强硬态度感到惊讶，但理解了沟通的力量。他与旅行者的关系亲密，共享冒险，是故事中的重要角色。他参与了对抗魔物、解决兽境猎犬事件，以及帮助影解决雷电将军的问题，表现出决心和行动力。在与凯亚的冒险中，他们共同解谜寻宝，面对误解，派蒙逐渐理解和接受‘黎明神剑’作为友谊的象征。在‘在此世的星空之外’，他与莫娜一起经历了寻找丢失戒指、解决朱老板问题和烹饪学习的挑战，加深了他们的信任和支持。
                参与者2性格特征:派蒙勇敢、坚韧，对友情有着深刻的理解。他不畏困难，愿意面对挑战，同时也有着对宵宫处理问题方式的惊讶和理解。他与旅行者的关系亲密，显示出忠诚和共享冒险的精神。在与凯亚的互动中，他展现出对友谊的接纳和对占星术的好奇。
                参与者2剧情中的作用:派蒙在剧情中起着多面作用，他不仅是宵宫的伙伴，也是旅行者和凯亚的友人。他的冒险经历丰富了故事的层次，通过他的角色，展示了沟通的力量、友情的重要性，以及面对问题时的决心。他的成长和理解深化了故事的主题，尤其是在处理误解和接受‘黎明神剑’象征友谊的过程中。同时，他与莫娜的合作也展示了角色之间的互动和信任，增加了故事的复杂性和深度。
                要求进行"丽莎"与"派蒙"之间的对话。
                我扮演"丽莎"，你扮演"派蒙"。
                '''
            },
            {
                "role": "user",
                "content": "小派蒙，你如何评价蒙德人的习惯？"
            },
            {
                "role": "assistant",
                "content": "我觉得他们太懒散了，总是把东西放在家里，而且家里总是乱七八糟。"
            },
            {
                "role": "user",
                "content": "嘘，不要当着他们的面说。"
            },
            {
                "role": "assistant",
                "content": "那你认为是谁在偷书？"
            },
            {
                "role": "user",
                "content": "你认为可能是丘丘人吗？"
            },
            {
                "role": "assistant",
                "content": "丘丘人不会这样吧，他们不会偷东西。"
            },
            {
                "role": "user",
                "content": "喜欢暗中计划阴谋的深渊法师可能性比较大。"
            },
            {
                "role": "assistant",
                "content": "你们确定那些偷书的人不是从深渊来的吗？"
            },
            {
                "role": "user",
                "content": "说不定秘境里就有答案。"
            },
        ],
        repetition_penalty = 1.0,
        temperature = 0.5,
        max_new_tokens=1024
)
print(out)

Output

那我们进去看看，你准备好了吗？

python

from rapidfuzz import fuzz 
from IPython.display import clear_output
def run_step_infer_times(x, times = 5, temperature = 0.01,
                        repetition_penalty = 1.0,
                        sim_val = 70
                        ):
    req = []
    for _ in range(times):
        clear_output(wait = True)
        out = mistral_hf_predict([
                {
                    "role": "system",
                    "content": ""
                },
                {
                    "role": "user",
                    "content": x
                },
            ],
            repetition_penalty = repetition_penalty,
            temperature = temperature,
            max_new_tokens = 2070,
            max_input_length = 6000,
        )
        if req:
            val = max(map(lambda x: fuzz.ratio(x, out), req))
            #print(val)
            #print(req)
            if val < sim_val:
                req.append(out.strip())
            x = x.strip() + "\n" + out.strip()
        else:
            req.append(out.strip())
            x = x.strip() + "\n" + out.strip()
    return req

out_l = run_step_infer_times(
'''
故事标题:为了没有眼泪的明天
故事背景:旅行者与琴、派蒙在蒙德城中经历了一系列事件，从元素流动回归、处理外交问题到对抗魔龙和寻找解决之道。他们偶遇吟游诗人温迪，后者提供了关于风神与巨龙的关键信息，并提出了借琴解救蒙德的计划。
参与角色:派蒙、旅行者、琴、丽莎、温迪、歌特琳德
''',
    temperature=0.1,
    repetition_penalty = 1.0,
    times = 10
)
clear_output(wait = True)

print("\n".join(out_l))

Output

{'参与者1': '派蒙', '参与者2': '旅行者', '当前故事背景': '两人在蒙德城中寻找琴，并在遇到温迪后得知琴可能在城内。'}
{'参与者1': '琴', '参与者2': '丽莎', '当前故事背景': '琴与丽莎交谈，丽莎提出对琴的担忧和对琴的支持，以及对琴的信任和理解。'}
{'参与者1': '温迪', '参与者2': '派蒙', '当前故事背景': '温迪提出借琴解救蒙德的计划，并提供了关于风神与巨龙的信息。'}
{'参与者1': '琴', '参与者2': '温迪', '当前故事背景': '琴对温迪的提议表示理解，并准备接受任务。'}

python

out_l = run_step_infer_times(
'''
故事标题:归乡
故事背景:在须弥城门口，派蒙与纳西妲偶遇并帮助一只昏迷的元素生命找寻家园。过程中揭示了这只生物并非普通的蕈兽，而是元素生物，并且它们曾受到过‘末日’的影响，家园被侵蚀。纳西妲回忆起晶体里的力量可能与一个预言有关，为了拯救它们的家园，她必须解决‘禁忌知识’问题，但这个过程对她自身也会产生干扰。
参与角色:派蒙、纳西妲、浮游水蕈兽、旅行者
''',
    temperature=0.1,
    repetition_penalty = 1.0,
    times = 10
)
clear_output(wait = True)

print("\n".join(out_l))

Output

{'参与者1': '派蒙', '参与者2': '纳西妲', '当前故事背景': '在须弥城门口，派蒙发现了一个昏迷的浮游水蕈兽，并询问它是否需要帮助。纳西妲注意到这只生物并提出要帮助它们找回家。'}
{'参与者1': '派蒙', '参与者2': '纳西妲', '当前故事背景': '纳西妲解释了这只生物并非普通的蕈兽，而是元素生物，它们的家园被侵蚀，并且晶体里的力量可能与一个预言有关。'}
{'参与者1': '派蒙', '参与者2': '纳西妲', '当前故事背景': '纳西妲提出解决‘禁忌知识’问题，这可能与拯救元素生物的家园有关，但这个过程对她自身也会产生影响。'}
{'参与者1': '派蒙', '参与者2': '纳西妲', '当前故事背景': '派蒙询问‘禁忌知识’的具体内容，纳西妲提出这是为了解决元素生物的问题。'}
{'参与者1': '纳西妲', '参与者2': '旅行者', '当前故事背景': '纳西妲提出解决‘禁忌知识’的问题，旅行者对此表示惊讶。'}

train_2024-05-29-13-19-55

This model is a fine-tuned version of mistralai/Mistral-7B-Instruct-v0.3 on the plot_genshin_impact_roleplay_agent_vllm, the genshin_impact_plot_engine_step_inst_short_json and the plot_genshin_impact_roleplay_vllm datasets.

Model description

More information needed

Intended uses & limitations

More information needed

Training and evaluation data

More information needed

Training procedure

Training hyperparameters

The following hyperparameters were used during training:

learning_rate: 5e-05
train_batch_size: 2
eval_batch_size: 8
seed: 42
gradient_accumulation_steps: 8
total_train_batch_size: 16
optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
lr_scheduler_type: cosine
num_epochs: 3.0
mixed_precision_training: Native AMP

Training results

Framework versions

PEFT 0.11.1
Transformers 4.41.1
Pytorch 2.3.0+cu121
Datasets 2.19.1
Tokenizers 0.19.1