alpaca-lora/utils/lima_prompter.py

"""
A dedicated helper to manage templates and prompt building.
"""

import json
import os.path as osp
from typing import List


class LimaPrompter(object):
    __slots__ = ("template", "_verbose")

    def __init__(self, template_name: str = "", verbose: bool = False):
        self._verbose = verbose
        if not template_name:
            template_name = "Lima"

    def generate_prompt(
        self,
        conversations: List[str]
    ) -> str:
        c = []
        for turn, text in enumerate(conversations):
            role = 'User' if turn % 2 == 0 else 'Assistant'
            c.append(f'### {role}: {text}')
        res = '\n'.join(c)
        if self._verbose:
            print(res)
        return res

    def get_response(self, output: str) -> str:
        res = output.split('### Assistant:')[-1].strip()
        if self._verbose:
            print(res)
        return res