from transformers import AutoModelForCausalLM, AutoTokenizer

checkpoint = "R0k1e/UltraLink-LM"

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
ultralink_lm = AutoModelForCausalLM.from_pretrained(checkpoint)

# Chat abilities in Chinese
# What is heavy cavalry?
first_question = "<s>[INST] 什么是重骑兵？ [/INST]"
chat_inputs = tokenizer.encode(first_question, add_special_tokens=False, return_tensors="pt")
chat_outputs = ultralink_lm.generate(chat_inputs, max_new_tokens=512)
first_answer = tokenizer.decode(chat_outputs[0])
print(first_answer)
# Expected output: 
"""
<s> [INST] 什么是重骑兵？ [/INST] 重骑兵是一种历史上的战斗单位，通常由骑兵组成，他们在战斗中使用重型装甲和长矛。他们以在战场上的强大攻击能力而闻名，并且通常被用于突破敌军阵线或攻击敌方骑兵。重骑兵通常被认为是中世纪战争中最强大和最具威慑力的单位之一。</s>
"""
# Translations in English:
"""
<s> [INST] What is heavy cavalry? [/INST] The Heavy Cavalry is a historical fighting unit, usually consisting of cavalrymen who use heavy armour and spears in battle. They were known for their powerful attacks on the battlefield and were often used to break through enemy lines or attack enemy cavalry. Heavy cavalry is often considered one of the most powerful and intimidating units in medieval warfare.</s>
"""

# Second turn:
second_question = "<s>[INST] What impact did heavy cavalry have on the medieval battlefield? [/INST]"
second_question = first_answer + second_question 
chat_inputs = tokenizer.encode(second_question, add_special_tokens=False, return_tensors="pt")
chat_outputs = ultralink_lm.generate(chat_inputs, max_new_tokens=512)
second_answer = tokenizer.decode(chat_outputs[0])
print(second_answer)
# Expected output: 
"""
<s> [INST] What is heavy cavalry? [/INST] The Heavy Cavalry is a historical fighting unit, usually consisting of cavalrymen who use heavy armour and spears in battle. They were known for their powerful attacks on the battlefield and were often used to break through enemy lines or attack enemy cavalry. Heavy cavalry is often considered one of the most powerful and intimidating units in medieval warfare.</s><s>[INST] What impact did heavy cavalry have on the medieval battlefield? [/INST] Heavy cavalry played an important role on the battlefields of the Middle Ages. Their powerful attacks made them a formidable force in battle and their presence could often change the outcome of a battle. Heavy cavalry were often used to break through enemy lines and create breaches in battle, which allowed other troops to move into enemy positions. They were also used to attack enemy cavalry, and their heavy armour and spears gave them an advantage in such situations. Overall, the presence of heavy cavalry made the medieval battlefield more complex and unpredictable, and their presence had a major impact on the outcome of battles.</s>
"""

# Code abilities in Russian
# Please implement a bubble sort algorithm in Python.
code_inputs = tokenizer.encode("<s>[INST] Реализуйте алгоритм пузырьковой сортировки на Python. [/INST]", add_special_tokens=False, return_tensors="pt")
code_outputs = ultralink_lm.generate(code_inputs, max_new_tokens=512)
print(tokenizer.decode(code_outputs[0]))
# Expected output: 
"""                                                                                 
<s> [INST] Реализуйте алгоритм пузырьковой сортировки на Python. [/INST] Конечно, вот простая реализация алгоритма пузырьковой сортировки на Python:

```python
def bubbleSort(arr):
    n = len(arr)

    # Проходимся по всем элементам массива
    for i in range(n):
        # Последние i элементов уже отсортированы
        for j in range(0, n-i-1):
            # Проходимся по массиву, пока не найдем элемент, больший, чем следующий
            if arr[j] > arr[j+1]:
                # Обмениваем элементы местами
                arr[j], arr[j+1] = arr[j+1], arr[j]

# Проверяем функцию
arr = [64, 34, 25, 12, 22, 11, 90]
bubbleSort(arr)
print("Отсортированный массив:", arr)
```

Эта реализация работает путем повторного прохождения по массиву и сравнения соседних элементов. Если элемент сравнивается больше, чем следующий элемент, они меняются местами. Этот процесс повторяется до тех пор, пока массив не будет отсортирован.

Вы можете изменить размер массива и его начальное состояние, чтобы увидеть, как функция работает с разными входными данными.

Примечание: Алгоритм пузырьковой сортировки не является эффективным для больших наборов данных, так как его средняя и худшая временная сложность составляют O(n^2), где n - количество элементов в массиве.</s>
"""
# Translations in English:
"""
<s> [INST] Implement the bubble sorting algorithm in Python. [/INST] Sure enough, here's a simple implementation of the bubble sort algorithm in Python:

```python
def bubbleSort(arr):
    n = len(arr)

    # Go through all elements of the array
    for i in range(n):
        # The last i elements are already sorted
        for j in range(0, n-i-1):
            # We traverse the array until we find an element greater than the next one
            if arr[j] > arr[j+1]:
                # Swapping elements
                arr[j], arr[j+1] = arr[j+1], arr[j]

# Check the function
arr = [64, 34, 25, 12, 22, 11, 90]
bubbleSort(arr)
print("Sorted array:", arr)
```

This implementation works by repeatedly traversing the array and comparing neighbouring elements. If an element is compared more than the next element, they are swapped. This process is repeated until the array is sorted.

You can change the array size and initial state to see how the function works with different input data.

Note: The bubble sort algorithm is not efficient for large datasets because its average and worst-case time complexity are O(n^2), where n is the number of elements in the array.</s>
"""

# Math abilities in French
# When the length of a rectangle is twice its width, calculate the area of the rectangle if its perimeter is 18 units?
math_inputs = tokenizer.encode("<s>[INST] Lorsque la longueur d'un rectangle est le double de sa largeur, calculer l'aire du rectangle si son périmètre est de 18 unités? [/INST]", add_special_tokens=False, return_tensors="pt")
math_outputs = ultralink_lm.generate(math_inputs, max_new_tokens=512)
print(tokenizer.decode(math_outputs[0]))
# Expected output: 
"""
<s> [INST] Lorsque la longueur d'un rectangle est le double de sa largeur, calculer l'aire du rectangle si son périmètre est de 18 unités? [/INST] 
Soit la largeur du rectangle $w$. Alors la longueur du rectangle est $2w$.
Le périmètre du rectangle est $2(w+2w)=18$.
En simplifiant, nous avons $6w=18$, donc $w=3$.
L'aire du rectangle est $w \cdot (2w) = 3 \cdot 6 = \boxed{18}$ unités carrées.
La réponse est : 18</s>
"""
# Translations in English:
"""
<s> [INST] When the length of a rectangle is twice its width, calculate the area of the rectangle if its perimeter is 18 units? [/INST] 
Let $w$ be the width of the rectangle. Then the length of the rectangle is $2w$.
La réponse est : 18
The perimeter of the rectangle is $2(w+2w)=18$. 
Simplifying, we have $6w=18$, so $w=3$. 
The area of the rectangle is $w \cdot (2w) = 3 \cdot 6 = \boxed{18}$ square units. 
The answer is: 18</s>
"""