File size: 5,549 Bytes
0ea5b31
 
 
 
134efaa
0ea5b31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
import streamlit as st
import numpy as np
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from pydantic.v1 import BaseModel, Field
from typing import Any, Optional, Dict, List
from huggingface_hub import InferenceClient
from langchain.llms.base import LLM
from markup import app_intro
import os

HF_token = os.getenv("apiToken")

model_name = "mistralai/Mistral-7B-Instruct-v0.1"
hf_token = HF_token
kwargs = {"max_new_tokens":10, "temperature":0.1, "top_p":0.95, "repetition_penalty":1.0, "do_sample":True}

class KwArgsModel(BaseModel):
    kwargs: Dict[str, Any] = Field(default_factory=dict)

class CustomInferenceClient(LLM, KwArgsModel):
    model_name: str
    inference_client: InferenceClient

    def __init__(self, model_name: str, hf_token: str, kwargs: Optional[Dict[str, Any]] = None):
        inference_client = InferenceClient(model=model_name, token=hf_token)
        super().__init__(
            model_name=model_name,
            hf_token=hf_token,
            kwargs=kwargs,
            inference_client=inference_client
        )

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None
    ) -> str:
        if stop is not None:
            raise ValueError("stop kwargs are not permitted.")
        response_gen = self.inference_client.text_generation(prompt, **self.kwargs, stream=True, return_full_text=False)
        response = ''.join(response_gen)  
        return response

    @property
    def _llm_type(self) -> str:
        return "custom"

    @property
    def _identifying_params(self) -> dict:
        return {"model_name": self.model_name}

def check_winner(board):
    for row in board:
        if len(set(row)) == 1 and row[0] != "":
            return row[0]
    for col in board.T:
        if len(set(col)) == 1 and col[0] != "":
            return col[0]
    if len(set(board.diagonal())) == 1 and board[0, 0] != "":
        return board[0, 0]
    if len(set(np.fliplr(board).diagonal())) == 1 and board[0, 2] != "":
        return board[0, 2]
    return None

def check_draw(board):
    return not np.any(board == "")

def main():
    st.set_page_config(page_title="Tic Tac Toe", page_icon=":memo:", layout="wide")

    col1, col2 = st.columns([1, 2])
    with col1:
        st.image("image.jpg", use_column_width=True)
    with col2:
        st.markdown(app_intro(), unsafe_allow_html=True) 

    st.markdown("____") 

    scores = st.session_state.get("scores", {"X": 0, "O": 0})
    board = np.array(st.session_state.get("board", [["" for _ in range(3)] for _ in range(3)]))
    current_player = st.session_state.get("current_player", "X")
    winner = check_winner(board)

    if winner is not None:
        scores[winner] += 1
        st.write(f"Player {winner} wins! Score: X - {scores['X']} | O - {scores['O']}")
    elif check_draw(board):
        st.write("Draw!")
    else:
        for row in range(3):
            cols = st.columns(3)
            for col in range(3):
                button_key = f"button_{row}_{col}"
                if board[row, col] == "" and current_player == "X":
                    if cols[col].button(" ", key=button_key):
                        board[row, col] = current_player
                        st.session_state.board = board
                        st.session_state.current_player = "O"
                        progress = st.session_state.get("progress", [])
                        progress.append(f"{current_player}: {chr(65 + row)}{col + 1}")
                        st.session_state.progress = progress
                        st.experimental_rerun()
                else:
                    cols[col].write(board[row, col])

    if current_player == "O" and winner is None:
        with st.spinner("Calculating AI Move..."):
            ai_progress = ", ".join(st.session_state.progress)
            ai_move = get_ai_move(ai_progress)
            ai_row, ai_col = ai_move.split(": ")[1]
            ai_row = ord(ai_row[0]) - 65
            ai_col = int(ai_col) - 1
            board[ai_row, ai_col] = "O"
            st.session_state.board = board

            progress = st.session_state.get("progress", [])
            progress.append(f"O: {chr(65 + ai_row)}{ai_col + 1}")
            st.session_state.progress = progress

            st.session_state.current_player = "X"
            st.experimental_rerun()

    st.markdown("____")
    if st.button("Reset game"):
        st.session_state.board = [["" for _ in range(3)] for _ in range(3)]
        st.session_state.current_player = "X"
        st.session_state.progress = []       
        st.experimental_rerun()

    if st.button("Reset Scores"):
        scores["X"] = 0
        scores["O"] = 0
        st.session_state.scores = scores
        st.experimental_rerun()

    
    progress = st.session_state.get("progress", [])
    st.write(", ".join(progress))

    st.write(f"Score: X - {scores['X']} | O - {scores['O']}")

def get_ai_move(progress):
    print("progress", progress)
    llm = CustomInferenceClient(model_name=model_name, hf_token=hf_token, kwargs=kwargs)

    template = """<s>[INST] Decide the next O move in tic tac toe game[/INST]
    Example output format = O: B2
    
    {progress}
    
    Next move:"""

    prompt = PromptTemplate(template=template, input_variables=["progress"])
    llm_chain = LLMChain(prompt=prompt, llm=llm)

    answer = llm_chain.run(progress)
    answer = answer.replace("</s>", "")
    
    print("ai move", answer)
    return answer


if __name__ == "__main__":
    main()