|
|
import asyncio |
|
|
import os |
|
|
from typing import Optional |
|
|
|
|
|
import gradio_client as grc |
|
|
from gradio_client.utils import QueueError |
|
|
|
|
|
|
|
|
HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
|
|
|
|
|
|
BOT_USER_ID = 1086256910572986469 if os.getenv("TEST_ENV", False) else 1102236653545861151 |
|
|
FALCON_CHANNEL_ID = 1079459939405279232 if os.getenv("TEST_ENV", False) else 1119313248056004729 |
|
|
|
|
|
|
|
|
thread_to_client = {} |
|
|
thread_to_user = {} |
|
|
|
|
|
|
|
|
async def wait(job): |
|
|
while not job.done(): |
|
|
await asyncio.sleep(0.2) |
|
|
|
|
|
|
|
|
def get_client(session: Optional[str] = None) -> grc.Client: |
|
|
client = grc.Client("huggingface-projects/falcon-180b-discord", hf_token=os.getenv("HF_TOKEN")) |
|
|
if session: |
|
|
client.session_hash = session |
|
|
return client |
|
|
|
|
|
|
|
|
async def falcon_chat(ctx, prompt): |
|
|
"""Generates text based on a given prompt""" |
|
|
try: |
|
|
if ctx.author.id != BOT_USER_ID: |
|
|
if ctx.channel.id == FALCON_CHANNEL_ID: |
|
|
if os.environ.get("TEST_ENV") == "True": |
|
|
print("Safetychecks passed for try_falcon") |
|
|
message = await ctx.send(f"**{prompt}** - {ctx.author.mention}") |
|
|
if len(prompt) > 99: |
|
|
small_prompt = prompt[:99] |
|
|
else: |
|
|
small_prompt = prompt |
|
|
thread = await message.create_thread(name=small_prompt, auto_archive_duration=60) |
|
|
|
|
|
if os.environ.get("TEST_ENV") == "True": |
|
|
print("Running falcon_initial_generation...") |
|
|
loop = asyncio.get_running_loop() |
|
|
client = await loop.run_in_executor(None, get_client, None) |
|
|
job = client.submit(prompt, api_name="/chat") |
|
|
await wait(job) |
|
|
try: |
|
|
job.result() |
|
|
response = job.outputs()[-1] |
|
|
await thread.send(response) |
|
|
thread_to_client[thread.id] = client |
|
|
thread_to_user[thread.id] = ctx.author.id |
|
|
except QueueError: |
|
|
await thread.send("The gradio space powering this bot is really busy! Please try again later!") |
|
|
except Exception as e: |
|
|
print(f"chat (180B) Error: {e}") |
|
|
|
|
|
|
|
|
async def continue_chat(message): |
|
|
"""Continues a given conversation based on chathistory""" |
|
|
try: |
|
|
if message.channel.id in thread_to_user: |
|
|
if thread_to_user[message.channel.id] == message.author.id: |
|
|
client = thread_to_client[message.channel.id] |
|
|
job = client.submit(message.content, api_name="/chat") |
|
|
await wait(job) |
|
|
try: |
|
|
job.result() |
|
|
response = job.outputs()[-1] |
|
|
await message.reply(response) |
|
|
except QueueError: |
|
|
await message.reply("The gradio space powering this bot is really busy! Please try again later!") |
|
|
except Exception as e: |
|
|
print(f"continue_falcon Error: {e}") |
|
|
await message.reply(f"Error: {e} <@811235357663297546> (continue_falcon error)") |
|
|
|
