Update README.md

README.md

@@ -28,10 +28,56 @@ import openai
 openai.api_base = "http://localhost:8000/v1"
 openai.api_key = "none"
 
+DEFAULT_MESSAGE_TEMPLATE = "<s>{role}\n{content}</s>"
+DEFAULT_RESPONSE_TEMPLATE = "<s>bot\n"
+DEFAULT_SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
+
+
+class Conversation:
+    def __init__(
+        self,
+        message_template=DEFAULT_MESSAGE_TEMPLATE,
+        system_prompt=DEFAULT_SYSTEM_PROMPT,
+        response_template=DEFAULT_RESPONSE_TEMPLATE
+    ):
+        self.message_template = message_template
+        self.response_template = response_template
+        self.messages = [{
+            "role": "system",
+            "content": system_prompt
+        }]
+
+    def add_user_message(self, message):
+        self.messages.append({
+            "role": "user",
+            "content": message
+        })
+
+    def add_bot_message(self, message):
+        self.messages.append({
+            "role": "bot",
+            "content": message
+        })
+
+    def get_prompt(self):
+        final_text = ""
+        for message in self.messages:
+            message_text = self.message_template.format(**message)
+            final_text += message_text
+        final_text += DEFAULT_RESPONSE_TEMPLATE
+        return final_text.strip()
+
+
+query = "Сколько весит жираф?"
+conversation = Conversation()
+conversation.add_user_message(query)
+prompt = conversation.get_prompt()
+
 response = openai.ChatCompletion.create(
         model="Gaivoronsky/Mistral-7B-Saiga",
-        messages=[{"role": "user", "content": 'Сколько весит жираф?'}],
+        messages=[{"role": "user", "content": prompt}],
         max_tokens=512,
+        stop=['</s>']
         )
 response['choices'][0]['message']['content']
 ```