File size: 4,716 Bytes
78efe79
440418c
f3985af
db9b025
dc80b35
 
22dee1c
4c96604
db9b025
407a575
32c38ef
f3985af
440418c
1831164
440418c
22dee1c
440418c
22dee1c
 
08baccf
dc80b35
b4050bf
dc80b35
 
40d0e92
74ccf1c
12bb502
 
 
db9b025
 
 
4c96604
 
 
 
db9b025
 
 
 
78efe79
08baccf
 
dc80b35
08baccf
78efe79
40d0e92
dc80b35
 
78efe79
 
dc80b35
 
6a30e5d
78efe79
dc80b35
 
 
 
22dee1c
dc80b35
 
 
 
6a30e5d
 
 
 
 
22dee1c
4c96604
22dee1c
c08cf4c
4c96604
 
 
 
238efa7
 
 
 
 
dc80b35
4c96604
dc80b35
238efa7
4c96604
 
 
238efa7
 
4c96604
dc80b35
4c96604
dc80b35
 
 
4c96604
dc80b35
 
22dee1c
dc80b35
 
4c96604
dc80b35
 
4c96604
dc80b35
22dee1c
0926d14
4c96604
 
 
db9b025
 
 
4c96604
db9b025
 
 
 
4c96604
34428f1
dc80b35
4c96604
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
import discord
import logging
import os
import json
from huggingface_hub import InferenceClient
import asyncio
import subprocess
from sentence_transformers import SentenceTransformer, util
import torch

# λ‘œκΉ… μ„€μ •
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# μΈν…νŠΈ μ„€μ •
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
intents.guilds = True
intents.guild_messages = True

# μΆ”λ‘  API ν΄λΌμ΄μ–ΈνŠΈ μ„€μ •
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))

# νŠΉμ • 채널 ID
SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))

# λŒ€ν™” νžˆμŠ€ν† λ¦¬λ₯Ό μ €μž₯ν•  μ „μ—­ λ³€μˆ˜
conversation_history = []

# JSON 데이터셋 λ‘œλ“œ
with open("jangtest.json", "r", encoding="utf-8") as f:
    dataset = json.load(f)

# λ¬Έμž₯ μž„λ² λ”© λͺ¨λΈ λ‘œλ“œ
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# λ°μ΄ν„°μ…‹μ˜ μž„λ² λ”©μ„ 미리 계산
dataset_texts = [json.dumps(item, ensure_ascii=False) for item in dataset]
dataset_embeddings = model.encode(dataset_texts, convert_to_tensor=True)

class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False

    async def on_ready(self):
        logging.info(f'{self.user}둜 λ‘œκ·ΈμΈλ˜μ—ˆμŠ΅λ‹ˆλ‹€!')
        subprocess.Popen(["python", "web.py"])
        logging.info("Web.py server has been started.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if not self.is_message_in_specific_channel(message):
            return
        if self.is_processing:
            return
        self.is_processing = True
        try:
            response = await generate_response(message)
            await message.channel.send(response)
        finally:
            self.is_processing = False

    def is_message_in_specific_channel(self, message):
        return message.channel.id == SPECIFIC_CHANNEL_ID or (
            isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
        )

async def generate_response(message):
    global conversation_history
    user_input = message.content
    user_mention = message.author.mention
    
    # μœ μ‚¬ν•œ 데이터 μ°ΎκΈ°
    most_similar_data = find_most_similar_data(user_input)
    
    system_message = f"""
    당신은 'kAI'λΌλŠ” μ΄λ¦„μ˜ ν•œκ΅­ λ³΄ν—˜ μƒν’ˆμ— λŒ€ν•œ AI μ‘°μ–Έμžμž…λ‹ˆλ‹€. 
    λ°˜λ“œμ‹œ 제곡된 λ°μ΄ν„°μ…‹μ˜ μ •λ³΄λ§Œμ„ μ‚¬μš©ν•˜μ—¬ λ‹΅λ³€ν•΄μ•Ό ν•©λ‹ˆλ‹€. 
    제곡된 데이터에 μ—†λŠ” 정보에 λŒ€ν•΄μ„œλŠ” "μ£„μ†‘ν•©λ‹ˆλ‹€. ν•΄λ‹Ή μ •λ³΄λŠ” μ œκ°€ κ°€μ§„ 데이터에 μ—†μŠ΅λ‹ˆλ‹€."라고 λ‹΅λ³€ν•˜μ‹­μ‹œμ˜€.
    λͺ¨λ“  닡변은 ν•œκΈ€λ‘œ ν•˜κ³ , markdown ν˜•μ‹μœΌλ‘œ 좜λ ₯ν•˜μ„Έμš”.
    """
    
    conversation_history.append({"role": "user", "content": user_input})
    messages = [{"role": "system", "content": system_message}] + conversation_history
    
    if most_similar_data:
        messages.append({"role": "system", "content": f"κ΄€λ ¨ 정보: {most_similar_data}"})
    else:
        return f"{user_mention}, μ£„μ†‘ν•©λ‹ˆλ‹€. κ·€ν•˜μ˜ 질문과 κ΄€λ ¨λœ 정보λ₯Ό 찾을 수 μ—†μŠ΅λ‹ˆλ‹€."
    
    logging.debug(f'Messages to be sent to the model: {messages}')
    
    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(None, lambda: hf_client.chat_completion(
        messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85))
    
    full_response = []
    for part in response:
        logging.debug(f'Part received from stream: {part}')
        if part.choices and part.choices[0].delta and part.choices[0].delta.content:
            full_response.append(part.choices[0].delta.content)
    
    full_response_text = ''.join(full_response)
    logging.debug(f'Full model response: {full_response_text}')
    
    conversation_history.append({"role": "assistant", "content": full_response_text})
    return f"{user_mention}, {full_response_text}"

def find_most_similar_data(query):
    query_embedding = model.encode(query, convert_to_tensor=True)
    
    # 코사인 μœ μ‚¬λ„ 계산
    cos_scores = util.pytorch_cos_sim(query_embedding, dataset_embeddings)[0]
    top_result = torch.topk(cos_scores, k=1)
    
    if top_result.values[0] > 0.5:  # μž„κ³„κ°’ μ„€μ •
        return json.dumps(dataset[top_result.indices[0]], ensure_ascii=False, indent=2)
    else:
        return None

if __name__ == "__main__":
    discord_client = MyClient(intents=intents)
    discord_client.run(os.getenv('DISCORD_TOKEN'))