Update README.md

README.md

@@ -8,8 +8,129 @@ language:
 pipeline_tag: conversational
 ---
 
+# Saiga 7B trained, Russian LLaMA-based chatbot
+
+Based on [LLaMA 7B](https://huggingface.co/huggyllama/llama-7b).
+
+* This is an adapter-only version.
+
 Colab: [link](https://colab.research.google.com/drive/1IBh4FMJPOGZAkX7DYWnIKdav_ZcKatlP)
 
+Training code: [link](https://github.com/IlyaGusev/rulm/tree/master/self_instruct)
+
+```python
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+
+MODEL_NAME = "IlyaGusev/saiga_7b_lora"
+
+config = PeftConfig.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    load_in_8bit=True,
+    device_map="auto"
+)
+model = PeftModel.from_pretrained(model, MODEL_NAME)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+DEFAULT_MESSAGE_TEMPLATE = "<s>{role}\n{content}</s>\n"
+DEFAULT_SYSTEM_PROMPT = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
+
+class Conversation:
+    def __init__(
+        self,
+        message_template=DEFAULT_MESSAGE_TEMPLATE,
+        system_prompt=DEFAULT_SYSTEM_PROMPT,
+        role_mapping=None,
+        start_token_id=1,
+        bot_token_id=9225
+    ):
+        self.message_template = message_template
+        self.role_mapping = role_mapping or {}
+        self.start_token_id = start_token_id
+        self.bot_token_id = bot_token_id
+        self.messages = [{
+            "role": "system",
+            "content": system_prompt
+        }]
+
+    def get_start_token_id(self):
+        return self.start_token_id
+
+    def get_bot_token_id(self):
+        return self.bot_token_id
+
+    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, tokenizer):
+        final_text = ""
+        for message in self.messages:
+            message_text = self.message_template.format(**message)
+            final_text += message_text
+        final_text += tokenizer.decode([self.start_token_id, self.bot_token_id])
+        return final_text.strip()
+
+def generate(model, tokenizer, prompt, generation_config):
+    data = tokenizer(prompt, return_tensors="pt")
+    data = {k: v.to(model.device) for k, v in data.items()}
+    output_ids = model.generate(
+        **data,
+        generation_config=generation_config
+    )[0]
+    output_ids = output_ids[len(data["input_ids"][0]):]
+    output = tokenizer.decode(output_ids, skip_special_tokens=True)
+    return output.strip()
+
+from transformers import GenerationConfig
+
+
+generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
+print(generation_config)
+
+inputs = ["Почему трава зеленая?", "Сочини длинный рассказ, обязательно упоминая следующие объекты. Дано: Таня, мяч"]
+for inp in inputs:
+    conversation = Conversation()
+    conversation.add_user_message(inp)
+    prompt = conversation.get_prompt(tokenizer)
+
+    output = generate(model, tokenizer, prompt, generation_config)
+    print(inp)
+    print(output)
+    print()
+    print("==============================")
+    print()
+```
+
+Examples:
+```
+Почему трава зеленая? 
+Зелень в цветении растений является результатом фотосинтеза, процесса, при котором растительные клетки превращают углекислый газ в кислород и энергию. Когда растение производит больше энергии, чем она потребляет, оно выделяет ее в виде красного или желтого цвета листьев. Однако, когда растение не получает достаточно света и воды, оно может изменить свой цвет на зеленый, чтобы сохранить свои запасы энергии для будущих потребностей. Зелень также может быть связана с защитой от ультрафиолетовых лучей, которые могут повредить растения.
+```
+
+```
+Сочини длинный рассказ, обязательно упоминая следующие объекты. Дано: Таня, мяч
+Выход: Таня любила футбол с детства. Она играла на улице со своими друзьями и мечтала стать профессиональным футболистом. Когда ей исполнилось 12 лет, она начала тренироваться в футбольной школе. Она была очень талантливой и упорной, и вскоре ей удалось попасть в молодежную команду.
+Однажды в лесу жил маленький медвежонок по имени Таня. Она была очень игривой и любила играть со своими друзьями - медведями. Однажды она решила провести время на свежем воздухе и отправилась в лес, чтобы играть с мячом. 
+Мяч был ее лучшим другом, который всегда находился рядом. Он был невероятно упрямым и не давал себя взять за руку. Но Таня никогда не сдавалась и продолжала пытаться его вытянуть из-под коры. 
+В лесу Таня встретила своего старого знакомого - медведя по имени Макс. Он был очень сильным и мощным, но он также любил играть с мячом. Их игра была так же веселой и приключенческой, как и у других медведей. 
+Но однажды, когда они играли, Таня заметила что-то необычное. В лесу началось грозить шторм, и она стала беспокоиться за своих друзей. Она решила вернуться домой, чтобы защитить их от грозового облака. 
+Когда Таня вернулась домой, она увидела, что все медведи были уже дома. Они говорили о том, что шторм прошел мимо, и они все остались живы. Таня была очень рада, что ее друзья остались в безопасности, и она поняла, что без мяча она была бы не такая веселая. 
+Так что Таня решила, что она должна найти свой самый лучший друг и вернуться к своим друзьям. Она начала искать его во всех местах, где он мог быть. Наконец, она нашла его под корнем дерева, где он был спрятан. 
+Таня была очень рада, что её лучший друг находится здесь, и она взяла его в руки. Она понимала, что без него она была бы совсем другой. Так что она поцеловала мяч и сказала ему, что она никогда не забудет его. 
+Так что теперь Таня и ее друзья играют вместе, и они всегда будут держать мяч в руках, чтобы сохранить свою дружбу и веселую атмосферу в лесу.
+
+```
+
 v2:
 - revision bef2597753a8f93c10422bf900e4c5bf7e0d186e
 - wandb [link](https://wandb.ai/ilyagusev/rulm_self_instruct/runs/8p3nfjqv/overview)