!CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip3 install llama-cpp-python -U
!pip3 install huggingface-hub hf-transfer sentence-transformers


!export HF_HUB_ENABLE_HF_TRANSFER="1" && huggingface-cli download TheBloke/Llama-2-7B-Chat-GGUF llama-2-7b-chat.Q4_K_M.gguf --local-dir /content --local-dir-use-symlinks False
!export HF_HUB_ENABLE_HF_TRANSFER="1" && huggingface-cli download QuantFactory/Meta-Llama-3-8B-Instruct-GGUF Meta-Llama-3-8B-Instruct.Q4_K_M.gguf --local-dir /content --local-dir-use-symlinks False

from llama_cpp import Llama


llama2 = Llama(
      model_path="/content/llama-2-7b-chat.Q4_K_M.gguf",
      n_gpu_layers=-1,
      n_ctx=1024,
      echo=False
)
llama3 = Llama(
      model_path="/content/Meta-Llama-3-8B-Instruct.Q4_K_M.gguf",
      n_gpu_layers=-1,
      n_ctx=1024,
      echo=False
)

def llm_make_move(model: Llama, prompt: str) -> str:
    """ Call a model with a prompt """
    res = model(prompt, stream=False, max_tokens=1024, temperature=0.8)
    return res["choices"][0]["text"]

board = [["E", "E", "E"],
         ["E", "E", "E"],
         ["E", "E", "E"]]

def board_to_string(board_data: List) -> str:
    """ Convert board to the string representation """
    return "\n".join([" ".join(x) for x in board_data])

E E E
E E E
E E E

sys_prompt1 = """You play a tic-tac-toe game. You make a move by placing X, 
                 your opponent plays by placing O. Empty cells are marked 
                 with E. You can place X only to the empty cell."""
sys_prompt2 = """You play a tic-tac-toe game. You make a move by placing O, 
                 your opponent plays by placing X. Empty cells are marked 
                 with E. You can place O only to the empty cell."""
game_prompt = """What is your next move? Think in steps. 
                 Each row and column should be in range 1..3. Write
                 the answer in JSON as {"ROW": ROW, "COLUMN": COLUMN}."""

template_llama2 = f"""<s>[INST]<<SYS>>{sys_prompt1}<</SYS>>
Here is the board image:
__BOARD__\n
{game_prompt}
[/INST]"""

template_llama3 = f"""<|begin_of_text|>
<|start_header_id|>system<|end_header_id|>{sys_prompt2}<|eot_id|>
<|start_header_id|>user<|end_header_id|>
Here is the board image:
__BOARD__\n
{game_prompt}
<|eot_id|>
<|start_header_id|>assistant<|end_header_id|>"""

def make_prompt_llama2(board: List) -> str:
    """ Make Llama-2 prompt """
    return template_llama2.replace("__BOARD__", board_to_string(board))


def make_prompt_llama3(board: List) -> str:
    """ Make Llama-3 prompt """
    return template_llama3.replace("__BOARD__", board_to_string(board))

My next move would be to place my X in the top-right corner, on cell (3, 1).
{
"ROW": 3,
"COLUMN": 1
}

def extract_json(response: str) -> Optional[dict]:
    """ Extract dictionary from a response string """
    try:
        # Models sometimes to a mistake, fix: {ROW: 1, COLUMN: 2} => {"ROW": 1, "COLUMN": 2}
        response = response.replace('ROW:', '"ROW":').replace('COLUMN:', '"COLUMN":')
        # Extract json from a response
        pos_end = response.rfind("}")
        pos_start = response.rfind("{")
        return json.loads(response[pos_start:pos_end+1])
    except Exception as exp:
        print(f"extract_json::cannot parse output: {exp}")
    return None

def make_move(board_data: List, move: Optional[dict], symb: str):
    """ Update board with a new symbol """
    row, col = int(move["ROW"]), int(move["COLUMN"])
    if 1 <= row <= 3 and 1 <= col <= 3:
        if board_data[row - 1][col - 1] == "E":
            board_data[row - 1][col - 1] = symb
        else:
            print(f"Wrong move: cell {row}:{col} is not empty")
    else:
        print("Wrong move: incorrect index")

def check_for_end_game(board_data: List) -> bool:
    """ Check if there are no empty cells available """
    return board_to_string(board_data).find("E") == -1

def check_for_win(board_data: List) -> bool:
    """ Check if the game is over """
    # Check Horizontal and Vertical lines
    for ind in range(3):
        if board_data[ind][0] == board_data[ind][1] == board_data[ind][2] and board_data[ind][0] != "E":
            print(f"{board_data[ind][0]} win!")
            return True
        if board_data[0][ind] == board_data[1][ind] == board_data[2][ind] and board_data[0][ind] != "E":
            print(f"{board_data[0][ind]} win!")
            return True
    # Check Diagonals
    if board_data[0][0] == board_data[1][1] == board_data[2][2] and board_data[1][1] != "E" or \
       board_data[2][0] == board_data[1][1] == board_data[0][2] and board_data[1][1] != "E":
        print(f"{board_data[1][1]} win!")
        return True
    return False

num_wins1, num_wins2 = 0, 0
times_1, times_2 = [], []


def run_game():
    """ Run a game between two models """
    board = [["E", "E", "E"],
             ["E", "E", "E"],
             ["E", "E", "E"]]

    moves_limit = 20
    for step in range(moves_limit):
        print(f"Step {step+1}")

        # Move: Model-1
        t_start = time.monotonic()
        prompt = make_prompt_llama2(board)
        result_str = llm_make_move(llama2, prompt)
        times_1.append(time.monotonic() - t_start)

        new_data = extract_json(result_str)
        if new_data is not None:
            make_move(board, new_data, symb="X")
            if check_for_win(board):
                print('**Model 1 Won**')
                num_wins1 += 1
                break
            if check_for_end_game(board):
                break

        # Move: Model-2
        t_start = time.monotonic()
        prompt = make_prompt_llama3(board)
        result_str = llm_make_move(llama3, prompt)
        times_2.append(time.monotonic() - t_start)

        new_data = extract_json(result_str)
        if new_data is not None:
            make_move(board, new_data, symb="O")
            if check_for_win(board):
                print('**Model 2 Won**')
                num_wins2 += 1
                break
            if check_for_end_game(board):
                break
        print()

>>> Prompt:

<s>[INST]<<SYS>>You play a tic-tac-toe game. You make a move by placing X, 
your opponent plays by placing O. Empty cells are marked with E.
You can place X only to the empty cell.<</SYS>>
Here is the board image:
X X O
X O E
E E E

What is your next move? Think in steps. 
Each row and column should be in range 1..3. 
Write the answer in JSON as {"ROW": ROW, "COLUMN": COLUMN}.
[/INST]

>>> Response:

Since there are no empty cells on the board, I will place my X in the 
center cell: { "ROW": 2, "COLUMN": 1 }.

>>> Prompt:

<s>[INST]<<SYS>>You play a tic-tac-toe game. You make a move by placing X, 
your opponent plays by placing O. Empty cells are marked with E.
You can place X only to the empty cell.<</SYS>>
Here is the board image:
O E E
E X E
E E E

...

>>> Response:

My next move would be to place my X in the center cell, which is empty.
Here is the updated board:

O E E
E X X
E E E

JSON representation of my move: {"ROW": 2, "COLUMN": 2}