Spaces:
Sleeping
Sleeping
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')) |