import chaiverse as chai
from chaiverse.formatters import PromptFormatter

template = """
### Instruction:

Enter roleplay mode. You are {character name}.

{character bio goes here}

Example session #1:
```
{character name}: words words ...
User: a reply
```

Example session #2:
```
...
```

### Input:
{character name}: words words words
{user name}: words words more words
{character name}: ...
{user name}: ...

### Response:
{character name}:
""".strip()

class AlpacaFormatter(PromptFormatter):
    memory_template = "### Instruction:\n\nEnter roleplay mode. You are {bot_name}.\n\n{memory}\n\n"
    prompt_template = "Example session #1:\n```\n{prompt}\n```\n\n### Input:\n"
    bot_template = "{bot_name}: {message}\n"
    user_template = "{user_name}: {message}\n"
    response_template = "\n### Response:\n{bot_name}:"

# model_url = "roleplay4fun/base-v0.2"
# model_url = "thanhdaonguyen/thanhdaorp-2712"
model_url = "roleplay4fun/christmas-7b-v0.3"

submission_parameters = {
	"model_repo": model_url,
    # "model_name": "christmas-v03",
	"generation_params": {
		"temperature": 0.72,
        "top_p": 0.73,
        "top_k": 40,
        "stopping_words": ['\n', '</s>', '<|im_end|>', '###'],
        "presence_penalty": 0.3,
        "frequency_penalty": 0.7,
        "max_input_tokens": 1024,
        "best_of": 8
	},
  "formatter": AlpacaFormatter(),
}
submitter = chai.ModelSubmitter(verbose=True)
submission_id = submitter.submit(submission_parameters)