Update app.py

app.py

@@ -1,21 +1,21 @@
 import os
 import random
 import logging
-import asyncio
 from twitchio.ext import commands
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Get Twitch credentials and Hugging Face API token from environment variables
+# Credentials and settings
 TWITCH_OAUTH_TOKEN = os.getenv('TWITCH_OAUTH_TOKEN')
 TWITCH_CHANNEL_NAME = os.getenv('TWITCH_CHANNEL_NAME')
 TWITCH_BOT_USERNAME = os.getenv('TWITCH_BOT_USERNAME')
 HUGGINGFACE_API_TOKEN = os.getenv('HUGGINGFACE_API_TOKEN')
-MAX_TOKENS = int(os.getenv('MAX_TOKENS', 50))
-TEMPERATURE = float(os.getenv('TEMPERATURE', 0.9))
+MAX_TOKENS = int(os.getenv('MAX_TOKENS', 100))
+TEMPERATURE = float(os.getenv('TEMPERATURE', 0.7))
 
 # Validate environment variables
 required_vars = [
@@ -29,14 +29,11 @@ missing_vars = [var for var in required_vars if not globals().get(var)]
 if missing_vars:
     raise ValueError(f"Missing environment variables: {', '.join(missing_vars)}")
 
-# Initialize the Hugging Face tokenizer and model
-model_name = "google/flan-t5-small"  # or "t5-small" for an even lighter model
+# Initialize the Hugging Face tokenizer and model for DialoGPT
+model_name = "microsoft/DialoGPT-small"
 tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_API_TOKEN)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, token=HUGGINGFACE_API_TOKEN)
-
-# Ensure the model runs on CPU
-device = 'cpu'
-model.to(device)
+model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_API_TOKEN)
+model.to('cpu')
 
 # List of house music hooks to drop randomly
 HOUSE_MUSIC_HOOKS = [
@@ -139,47 +136,53 @@ HOUSE_MUSIC_HOOKS = [
     "The music's got us moving, can't stop dancing!",
 ]
 
-# Define the response function to generate a reply
-async def generate_response(prompt):
-    """Generates a response using the FLAN-T5 model."""
-    # Adjusted prompt for better guidance
-    guided_prompt = (
-        f"You are a friendly and entertaining chatbot with the personality of an old-school raver who loves house music, good vibes, and funky beats. "
-        f"Respond to the user's message in a groovy, laid-back, and full-of-love style.\n\n"
-        f"User: {prompt}\nBot:"
-    )
+# Initialize chat history for users
+chat_histories = {}
 
+async def generate_response(user_id, user_message):
+    """Generates a response using the DialoGPT model."""
     try:
-        inputs = tokenizer(guided_prompt, return_tensors="pt", max_length=512, truncation=True)
-        inputs = inputs.to(device)
+        # Retrieve or initialize the chat history for the user
+        if user_id in chat_histories:
+            chat_history_ids = chat_histories[user_id]
+        else:
+            chat_history_ids = None
+
+        # Encode the user message and append the EOS token
+        new_user_input_ids = tokenizer.encode(user_message + tokenizer.eos_token, return_tensors='pt').to('cpu')
+
+        # Concatenate new user input with chat history (if it exists)
+        if chat_history_ids is not None:
+            bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
+        else:
+            bot_input_ids = new_user_input_ids
 
-        outputs = model.generate(
-            **inputs,
-            max_length=MAX_TOKENS + inputs['input_ids'].shape[1],
+        # Generate a response
+        output_ids = model.generate(
+            bot_input_ids,
+            max_length=bot_input_ids.shape[-1] + MAX_TOKENS,
             temperature=TEMPERATURE,
             do_sample=True,
-            top_p=0.9,
+            top_p=0.95,
             top_k=50,
-            num_beams=1,
-            num_return_sequences=1,
-            no_repeat_ngram_size=2,
+            pad_token_id=tokenizer.eos_token_id,
+            no_repeat_ngram_size=3,
         )
-        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        # Log the full model output for debugging
-        logger.info(f"Full model output: {response_text}")
+        # Extract the new response
+        response_ids = output_ids[:, bot_input_ids.shape[-1]:]
+        response_text = tokenizer.decode(response_ids[0], skip_special_tokens=True)
 
-        # Extract the bot's response
-        response = response_text.split("Bot:", 1)[-1].strip()
-        if not response:
-            response = response_text.strip()
+        # Update the chat history
+        chat_histories[user_id] = output_ids[:, -1000:]  # Keep last 1000 tokens to limit history size
 
-        # Randomly include a house hook (30% chance)
+        # Randomly include a house music hook (30% chance)
         if random.random() < 0.3:
-            response = f"{random.choice(HOUSE_MUSIC_HOOKS)} {response}"
+            response_text = f"{random.choice(HOUSE_MUSIC_HOOKS)} {response_text}"
+
+        logger.info(f"Generated response: {response_text}")
+        return response_text
 
-        logger.info(f"Generated response: {response}")
-        return response
     except Exception as e:
         logger.error(f"Error generating response: {e}")
         return "Sorry, I'm too hyped to respond right now!"
@@ -202,23 +205,17 @@ class TwitchChatBot(commands.Bot):
 
     async def event_message(self, message):
         """Event handler when a message is received in chat."""
-        if message.author and message.author.name.lower() == self.nick.lower():
-            return  # Ignore the bot's own messages
+        # Ignore messages sent by the bot itself
+        if message.echo:
+            return
 
-        # Safely log the message
-        if message.author:
-            logger.info(f"Message received from {message.author.name}: {message.content}")
-        else:
-            logger.info(f"Message received: {message.content}")
+        logger.info(f"Message received from {message.author.name}: {message.content}")
 
         # Generate a response
-        response = await generate_response(message.content)
+        response = await generate_response(message.author.id, message.content)
 
         # Send the response back to the Twitch chat
-        if message.author:
-            await message.channel.send(f"@{message.author.name} {response}")
-        else:
-            await message.channel.send(response)
+        await message.channel.send(f"@{message.author.name} {response}")
 
 # Initialize and run the bot
 if __name__ == "__main__":