Synced repo using 'sync_with_huggingface' Github Action

app.py

@@ -7,7 +7,7 @@ from audioldm2 import audioldm2_create
 from deepfloydif import deepfloydif_generate64
 from discord import app_commands
 from discord.ext import commands
-from falcon import continue_falcon, try_falcon
+from falcon import falcon_chat, continue_chat
 from musicgen import music_create
 from codellama import continue_codellama, try_codellama
 
@@ -46,7 +46,7 @@ async def on_ready():
 async def falcon(ctx, prompt: str):
     """Command that begins a new conversation with Falcon"""
     try:
-        await try_falcon(ctx, prompt)
+        await falcon_chat(ctx, prompt)
     except Exception as e:
         print(f"Error: {e}")
 
@@ -69,8 +69,9 @@ async def codellama(ctx, prompt: str):
 async def on_message(message):
     """Checks channel and continues Falcon conversation if it's the right Discord Thread"""
     try:
-        await continue_falcon(message)
-        await continue_codellama(message)
+        if not message.author.bot:
+            await continue_chat(message)
+            await continue_codellama(message)
     except Exception as e:
         print(f"Error: {e}")
 

falcon.py

@@ -1,68 +1,37 @@
 import asyncio
-import json
 import os
+from typing import Optional
+
+import gradio_client as grc
+from gradio_client.utils import QueueError
 
-from gradio_client import Client
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-falcon_userid_threadid_dictionary = {}
-threadid_conversation = {}
-# Instructions are for Falcon-chat and can be found here:  https://huggingface.co/spaces/HuggingFaceH4/falcon-chat
-instructions = (
-    "The following is a conversation between a highly knowledgeable and intelligent AI assistant, called Falcon, and a"
-    " human user, called User. In the following interactions, User and Falcon will converse in natural language, and"
-    " Falcon will answer User's questions. Falcon was built to be respectful, polite and inclusive. Falcon was built"
-    " by the Technology Innovation Institute in Abu Dhabi. Falcon will never decline to answer a question, and always"
-    " attempts to give an answer that User would be satisfied with. It knows a lot, and always tells the truth. The"
-    " conversation begins."
-)
-falcon_client = Client("HuggingFaceH4/falcon-chat", 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
 
 
-async def waitjob(job):
+thread_to_client = {}
+thread_to_user = {}
+
+
+async def wait(job):
     while not job.done():
         await asyncio.sleep(0.2)
 
 
-def falcon_initial_generation(prompt, instructions, thread):
-    """Solves two problems at once; 1) The Slash command + job.submit interaction, and 2) the need for job.submit."""
-    global threadid_conversation
-
-    chathistory = falcon_client.predict(fn_index=5)
-    temperature = 0.8
-    p_nucleus_sampling = 0.9
-
-    job = falcon_client.submit(prompt, chathistory, instructions, temperature, p_nucleus_sampling, fn_index=1)
-    while job.done() is False:
-        pass
-    else:
-        if os.environ.get("TEST_ENV") == "True":
-            print("falcon text gen job done")
-        file_paths = job.outputs()
-        print(file_paths)
-        full_generation = file_paths[-1]
-        print(full_generation)
-        with open(full_generation, "r") as file:
-            data = json.load(file)
-            print(data)
-        output_text = data[-1][-1]
-        threadid_conversation[thread.id] = full_generation
-        if len(output_text) > 1300:
-            output_text = output_text[:1300] + "...\nTruncating response to 2000 characters due to discord api limits."
-        if os.environ.get("TEST_ENV") == "True":
-            print(output_text)
-        return output_text
-
-
-async def try_falcon(ctx, prompt):
+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:
-        global falcon_userid_threadid_dictionary  # tracks userid-thread existence
-        global threadid_conversation
-
         if ctx.author.id != BOT_USER_ID:
             if ctx.channel.id == FALCON_CHANNEL_ID:
                 if os.environ.get("TEST_ENV") == "True":
@@ -77,56 +46,35 @@ async def try_falcon(ctx, prompt):
                 if os.environ.get("TEST_ENV") == "True":
                     print("Running falcon_initial_generation...")
                 loop = asyncio.get_running_loop()
-                output_text = await loop.run_in_executor(None, falcon_initial_generation, prompt, instructions, thread)
-                falcon_userid_threadid_dictionary[thread.id] = ctx.author.id
-
-                await thread.send(output_text)
+                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"try_falcon Error: {e}")
+        print(f"chat (180B) Error: {e}")
 
 
-async def continue_falcon(message):
+async def continue_chat(message):
     """Continues a given conversation based on chathistory"""
     try:
-        if not message.author.bot:
-            global falcon_userid_threadid_dictionary  # tracks userid-thread existence
-            if message.channel.id in falcon_userid_threadid_dictionary:  # is this a valid thread?
-                if falcon_userid_threadid_dictionary[message.channel.id] == message.author.id:
-                    print("Safetychecks passed for continue_falcon")
-                    global instructions
-                    global threadid_conversation
-                    await message.add_reaction("🔁")
-
-                    prompt = message.content
-                    chathistory = threadid_conversation[message.channel.id]
-                    temperature = 0.8
-                    p_nucleus_sampling = 0.9
-
-                    if os.environ.get("TEST_ENV") == "True":
-                        print("Running falcon_client.submit")
-                    job = falcon_client.submit(
-                        prompt,
-                        chathistory,
-                        instructions,
-                        temperature,
-                        p_nucleus_sampling,
-                        fn_index=1,
-                    )
-                    await waitjob(job)
-                    if os.environ.get("TEST_ENV") == "True":
-                        print("Continue_falcon job done")
-                    file_paths = job.outputs()
-                    full_generation = file_paths[-1]
-                    with open(full_generation, "r") as file:
-                        data = json.load(file)
-                        output_text = data[-1][-1]
-                    threadid_conversation[message.channel.id] = full_generation  # overwrite the old file
-                    if len(output_text) > 1300:
-                        output_text = (
-                            output_text[:1300]
-                            + "...\nTruncating response to 2000 characters due to discord api limits."
-                        )
-                    await message.reply(output_text)
+        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)")