s

app.py

@@ -42,7 +42,7 @@ async def on_message(message):
         message_counts[message.channel.id] = 0 # Reset the counter
 
 
-    response = generate(str(message), temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0)
+    response = ai.text_generation(message, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0)
     splited_response = split_string(str(response))
     for part in splited_response:
         await channel.send(part)
@@ -53,30 +53,6 @@ def split_string(text: str) -> list[str]:
     """Helper function to split text into chunks"""
     return [text[i:i+3900] for i in range(0, len(text), 3900)]
 
-def generate(
-    prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0
-):
-    temperature = float(temperature)
-    if temperature < 1e-2:
-        temperature = 1e-2
-    top_p = float(top_p)
-
-    generate_kwargs = dict(
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=42,
-    )
-
-    stream = ai.text_generation(prompt, **generate_kwargs, stream=False, details=True, return_full_text=True)
-    output = ""
-    for response in stream:
-        output += response.token.text
-        yield output
-    return output
-
 @bot.event
 async def on_ready():
     print(f'Logged in as {bot.user}')  # Logs bot login in console