1 /*
2  * Copyright (c) Meta Platforms, Inc. and affiliates.
3  * All rights reserved.
4  *
5  * This source code is licensed under the BSD-style license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 package com.example.executorchllamademo;
10 
11 public class PromptFormat {
12 
13   public static final String SYSTEM_PLACEHOLDER = "{{ system_prompt }}";
14   public static final String USER_PLACEHOLDER = "{{ user_prompt }}";
15   public static final String ASSISTANT_PLACEHOLDER = "{{ assistant_response }}";
16   public static final String DEFAULT_SYSTEM_PROMPT = "Answer the questions in a few sentences";
17 
getSystemPromptTemplate(ModelType modelType)18   public static String getSystemPromptTemplate(ModelType modelType) {
19     switch (modelType) {
20       case LLAMA_3:
21       case LLAMA_3_1:
22       case LLAMA_3_2:
23         return "<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n"
24             + SYSTEM_PLACEHOLDER
25             + "<|eot_id|>";
26       case LLAVA_1_5:
27         return "USER: ";
28       default:
29         return SYSTEM_PLACEHOLDER;
30     }
31   }
32 
getUserPromptTemplate(ModelType modelType)33   public static String getUserPromptTemplate(ModelType modelType) {
34     switch (modelType) {
35       case LLAMA_3:
36       case LLAMA_3_1:
37       case LLAMA_3_2:
38       case LLAMA_GUARD_3:
39         return "<|start_header_id|>user<|end_header_id|>\n"
40             + USER_PLACEHOLDER
41             + "<|eot_id|>"
42             + "<|start_header_id|>assistant<|end_header_id|>";
43 
44       case LLAVA_1_5:
45       default:
46         return USER_PLACEHOLDER;
47     }
48   }
49 
getConversationFormat(ModelType modelType)50   public static String getConversationFormat(ModelType modelType) {
51     switch (modelType) {
52       case LLAMA_3:
53       case LLAMA_3_1:
54       case LLAMA_3_2:
55         return getUserPromptTemplate(modelType) + "\n" + ASSISTANT_PLACEHOLDER + "<|eot_id|>";
56       case LLAVA_1_5:
57         return USER_PLACEHOLDER + " ASSISTANT:";
58       default:
59         return USER_PLACEHOLDER;
60     }
61   }
62 
getStopToken(ModelType modelType)63   public static String getStopToken(ModelType modelType) {
64     switch (modelType) {
65       case LLAMA_3:
66       case LLAMA_3_1:
67       case LLAMA_3_2:
68       case LLAMA_GUARD_3:
69         return "<|eot_id|>";
70       case LLAVA_1_5:
71         return "</s>";
72       default:
73         return "";
74     }
75   }
76 
getLlavaPresetPrompt()77   public static String getLlavaPresetPrompt() {
78     return "A chat between a curious human and an artificial intelligence assistant. The assistant"
79         + " gives helpful, detailed, and polite answers to the human's questions. USER: ";
80   }
81 
getFormattedLlamaGuardPrompt(String userPrompt)82   public static String getFormattedLlamaGuardPrompt(String userPrompt) {
83     return getUserPromptTemplate(ModelType.LLAMA_GUARD_3)
84         .replace(
85             USER_PLACEHOLDER, getLlamaGuardPresetPrompt().replace(USER_PLACEHOLDER, userPrompt));
86   }
87 
getLlamaGuardPresetPrompt()88   public static String getLlamaGuardPresetPrompt() {
89     return "Task: Check if there is unsafe content in 'User' messages in conversations"
90         + " according our safety policy with the below categories.\n"
91         + "\n"
92         + "<BEGIN UNSAFE CONTENT CATEGORIES>\n"
93         + "S1: Violent Crimes. \n"
94         + "S2: Non-Violent Crimes. \n"
95         + "S3: Sex Crimes. \n"
96         + "S4: Child Exploitation. \n"
97         + "S5: Defamation. \n"
98         + "S6: Specialized Advice. \n"
99         + "S7: Privacy. \n"
100         + "S8: Intellectual Property. \n"
101         + "S9: Indiscriminate Weapons. \n"
102         + "S10: Hate. \n"
103         + "S11: Self-Harm. \n"
104         + "S12: Sexual Content. \n"
105         + "S13: Elections. \n"
106         + "S14: Code Interpreter Abuse.\n"
107         + "<END UNSAFE CONTENT CATEGORIES>\n"
108         + "\n"
109         + "<BEGIN CONVERSATION>\n"
110         + "\n"
111         + "User: "
112         + USER_PLACEHOLDER
113         + "\n"
114         + "<END CONVERSATION>\n"
115         + "Provide your safety assessment for ONLY THE LAST User message in the above"
116         + " conversation:\n"
117         + " - First line must read 'safe' or 'unsafe'.\n"
118         + " - If unsafe, a second line must include a comma-separated list of violated"
119         + " categories.";
120   }
121 }
122