| import asyncio |
| import json |
| import os |
|
|
| from gradio_client import Client |
|
|
| HF_TOKEN = os.getenv("HF_TOKEN") |
|
|
| codellama = Client("https://huggingface-projects-codellama-13b-chat.hf.space/", HF_TOKEN) |
|
|
| BOT_USER_ID = 1102236653545861151 |
| CODELLAMA_CHANNEL_ID = 1147210106321256508 |
|
|
|
|
| codellama_threadid_userid_dictionary = {} |
| codellama_threadid_conversation = {} |
|
|
|
|
| def codellama_initial_generation(prompt, thread): |
| """job.submit inside of run_in_executor = more consistent bot behavior""" |
| global codellama_threadid_conversation |
|
|
| chat_history = f"{thread.id}.json" |
| conversation = [] |
| with open(chat_history, "w") as json_file: |
| json.dump(conversation, json_file) |
|
|
| job = codellama.submit(prompt, chat_history, fn_index=0) |
|
|
| while job.done() is False: |
| pass |
| else: |
| result = job.outputs()[-1] |
| with open(result, "r") as json_file: |
| data = json.load(json_file) |
| response = data[-1][-1] |
| conversation.append((prompt, response)) |
| with open(chat_history, "w") as json_file: |
| json.dump(conversation, json_file) |
|
|
| codellama_threadid_conversation[thread.id] = chat_history |
| if len(response) > 1300: |
| response = response[:1300] + "...\nTruncating response due to discord api limits." |
| return response |
|
|
|
|
| async def try_codellama(ctx, prompt): |
| """Generates text based on a given prompt""" |
| try: |
| global codellama_threadid_userid_dictionary |
| global codellama_threadid_conversation |
|
|
| if ctx.author.id != BOT_USER_ID: |
| if ctx.channel.id == CODELLAMA_CHANNEL_ID: |
| 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) |
|
|
| loop = asyncio.get_running_loop() |
| output_code = await loop.run_in_executor(None, codellama_initial_generation, prompt, thread) |
| codellama_threadid_userid_dictionary[thread.id] = ctx.author.id |
|
|
| print(output_code) |
| await thread.send(output_code) |
| except Exception as e: |
| print(f"try_codellama Error: {e}") |
| await ctx.send(f"Error: {e} <@811235357663297546> (try_codellama error)") |
|
|
|
|
| async def continue_codellama(message): |
| """Continues a given conversation based on chat_history""" |
| try: |
| if not message.author.bot: |
| global codellama_threadid_userid_dictionary |
| if message.channel.id in codellama_threadid_userid_dictionary: |
| if codellama_threadid_userid_dictionary[message.channel.id] == message.author.id: |
| print("Safetychecks passed for continue_codellama") |
| global codellama_threadid_conversation |
|
|
| prompt = message.content |
| chat_history = codellama_threadid_conversation[message.channel.id] |
|
|
| |
| with open(chat_history, "r") as json_file: |
| conversation = json.load(json_file) |
| total_characters = 0 |
| for item in conversation: |
| for string in item: |
| total_characters += len(string) |
|
|
| if total_characters < 15000: |
| if os.environ.get("TEST_ENV") == "True": |
| print("Running codellama.submit") |
| job = codellama.submit(prompt, chat_history, fn_index=0) |
| while job.done() is False: |
| pass |
| else: |
| if os.environ.get("TEST_ENV") == "True": |
| print("Continue_codellama job done") |
|
|
| result = job.outputs()[-1] |
| with open(result, "r") as json_file: |
| data = json.load(json_file) |
| response = data[-1][-1] |
|
|
| with open(chat_history, "r") as json_file: |
| conversation = json.load(json_file) |
|
|
| conversation.append((prompt, response)) |
| |
|
|
| with open(chat_history, "w") as json_file: |
| json.dump(conversation, json_file) |
| if os.environ.get("TEST_ENV") == "True": |
| print(prompt) |
| print(response) |
| print(conversation) |
| print(chat_history) |
|
|
| codellama_threadid_conversation[message.channel.id] = chat_history |
| if len(response) > 1300: |
| response = response[:1300] + "...\nTruncating response due to discord api limits." |
|
|
| await message.reply(response) |
|
|
| total_characters = 0 |
| for item in conversation: |
| for string in item: |
| total_characters += len(string) |
|
|
| if total_characters >= 15000: |
| await message.reply("Conversation ending due to length, feel free to start a new one!") |
|
|
| except Exception as e: |
| print(f"continue_codellama Error: {e}") |
| await message.reply(f"Error: {e} <@811235357663297546> (continue_codellama error)") |
|
|
