Update app.py

app.py

@@ -1,140 +1,171 @@
-import os
-import random
-import logging
-from twitchio.ext import commands
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# 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', 100))
-TEMPERATURE = float(os.getenv('TEMPERATURE', 0.7))
-
-# Validate environment variables
-required_vars = [
-    'TWITCH_OAUTH_TOKEN',
-    'TWITCH_CHANNEL_NAME',
-    'TWITCH_BOT_USERNAME',
-    'HUGGINGFACE_API_TOKEN'
-]
-
-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 for TinyLlama
-model_name = "jacksonstrut/tinyllama-1.1B-chat"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HUGGINGFACE_API_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HUGGINGFACE_API_TOKEN)
-model.to('cpu')
-
-# List of house music hooks to drop randomly
-HOUSE_MUSIC_HOOKS = [
-    "Feeling the vibe, let's get hyped!",
-    "Rave or die!",
-    "Rave to the g,rave!",
-    "He just strutts!",
-    "Dancing is what to do!",
-    "Let's get riggity riggity wrecked!",
-    "Sleep is for the week, not the weekend!",
-    "It's hard out here for a pimp!",
-    "Something for your mind!",
-    "Something for your body!",
-    "Something for your soul!",
-    "Not everyone understands this thing they call house music!",
-    "It's a spiritual thing!",
-    # ... [keep the rest of the hooks as they are] ...
-]
-
-# Initialize chat history for users
-chat_histories = {}
-
-async def generate_response(user_id, user_message):
-    """Generates a response using the TinyLlama model."""
-    try:
-        # 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
-
-        # 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.95,
-            top_k=50,
-            pad_token_id=tokenizer.eos_token_id,
-            no_repeat_ngram_size=3,
-        )
-
-        # Extract the new response
-        response_ids = output_ids[:, bot_input_ids.shape[-1]:]
-        response_text = tokenizer.decode(response_ids[0], skip_special_tokens=True)
-
-        # Update the chat history
-        chat_histories[user_id] = output_ids[:, -1000:]  # Keep last 1000 tokens to limit history size
-
-        # Randomly include a house music hook (30% chance)
-        if random.random() < 0.3:
-            response_text = f"{random.choice(HOUSE_MUSIC_HOOKS)} {response_text}"
-
-        logger.info(f"Generated response: {response_text}")
-        return response_text
-
-    except Exception as e:
-        logger.error(f"Error generating response: {e}")
-        return "Sorry, I'm too hyped to respond right now!"
-
-# Create a Twitch chatbot using TwitchIO
-class TwitchChatBot(commands.Bot):
-
-    def __init__(self):
-        super().__init__(
-            token=TWITCH_OAUTH_TOKEN,
-            nick=TWITCH_BOT_USERNAME,
-            prefix='!',
-            initial_channels=[TWITCH_CHANNEL_NAME]
-        )
-
-    async def event_ready(self):
-        """Event handler when the bot is connected and ready."""
-        logger.info(f"Logged in as | {self.nick}")
-        logger.info(f"Connected to channel | {TWITCH_CHANNEL_NAME}")
-
-    async def event_message(self, message):
-        """Event handler when a message is received in chat."""
-        # Ignore messages sent by the bot itself
-        if message.echo:
-            return
-
-        logger.info(f"Message received from {message.author.name}: {message.content}")
-
-        # Generate a response
-        response = await generate_response(message.author.id, message.content)
-
-        # Send the response back to the Twitch chat
-        await message.channel.send(f"@{message.author.name} {response}")
-
-# Initialize and run the bot
-if __name__ == "__main__":
-    bot = TwitchChatBot()
-    bot.run()
+import torch.nn as nn
+import yaml
+from transformers import GPT2Config, PreTrainedTokenizerFast
+from huggingface_hub import HfApi, upload_file
+import os
+import json
+import sentencepiece as spm
+from tokenizers import SentencePieceBPETokenizer
+
+# Define the TinyLlama Model
+class TinyLlama(nn.Module):
+    def __init__(self, config):
+        super(TinyLlama, self).__init__()
+        self.embedding = nn.Embedding(config.vocab_size, config.n_embd)
+        self.transformer_blocks = nn.ModuleList([
+            nn.TransformerEncoderLayer(
+                d_model=config.n_embd,
+                nhead=config.n_head,
+                dim_feedforward=config.hidden_dim
+            )
+            for _ in range(config.n_layer)
+        ])
+        self.output_layer = nn.Linear(config.n_embd, config.vocab_size)
+
+    def forward(self, x):
+        x = self.embedding(x)
+        for block in self.transformer_blocks:
+            x = block(x)
+        x = self.output_layer(x)
+        return x
+
+# Load Model Configuration from YAML file
+model_config_path = "/home/jax/out/custom-model/final/model_config.yaml"
+with open(model_config_path, 'r') as file:
+    config_data = yaml.safe_load(file)
+
+# Create Model Configuration
+config = GPT2Config(
+    vocab_size=config_data.get("vocab_size", 32000),    # Adjust to match TinyLlama's vocab size
+    n_embd=config_data.get("n_embd", 2048),             # Embedding size for TinyLlama
+    n_layer=config_data.get("n_layer", 24),             # Number of transformer layers
+    n_head=config_data.get("n_head", 16),               # Number of attention heads
+    hidden_dim=config_data.get("hidden_dim", 8192)      # Feedforward layer dimension
+)
+
+# Instantiate the TinyLlama Model
+model = TinyLlama(config)
+
+# Load Weights from the .pth file
+model_weights_path = "/home/jax/out/custom-model/final/lit_model.pth"
+model_weights = torch.load(model_weights_path, map_location=torch.device('cpu'), weights_only=True)
+model.load_state_dict(model_weights, strict=False)  # strict=False to allow for minor mismatches
+
+# Directory to Save the Model and Tokenizer for Hugging Face
+model_dir = "./huggingface_tinyllama"
+os.makedirs(model_dir, exist_ok=True)
+
+# Save the Model Weights
+model_weights_save_path = os.path.join(model_dir, "pytorch_model.bin")
+torch.save(model.state_dict(), model_weights_save_path)
+
+# Save the Configuration in JSON format
+config_save_path = os.path.join(model_dir, "config.json")
+with open(config_save_path, 'w') as f:
+    json.dump(config.to_dict(), f)
+
+# Load and Convert the SentencePiece model to Hugging Face-compatible format
+# Load the SentencePiece model
+tokenizer_path = "/home/jax/out/custom-model/final/tokenizer.model"
+sp_tokenizer = spm.SentencePieceProcessor()
+sp_tokenizer.load(tokenizer_path)
+
+# Create a new tokenizer using the tokenizers library
+hf_tokenizer = SentencePieceBPETokenizer(
+    vocab=tokenizer_path,  # Load SentencePiece model
+    add_prefix_space=True
+)
+
+# Train the tokenizer if needed (optional step, depending on your model and usage)
+# hf_tokenizer.train([tokenizer_path])
+
+# Set special tokens
+hf_tokenizer.add_special_tokens(["<unk>", "<pad>", "<s>", "</s>"])
+
+# Wrap the tokenizer with PreTrainedTokenizerFast for Hugging Face compatibility
+hf_tokenizer = PreTrainedTokenizerFast.from_pretrained(tokenizer_path)
+
+# Save the Hugging Face tokenizer
+hf_tokenizer.save_pretrained(model_dir)
+
+# Push to Hugging Face Hub
+api.create_repo(repo_id=repo_id, token=os.getenv('HUGGINGFACE_API_TOKEN'), exist_ok=True)
+
+# Create a Model Card for the Model
+model_card = """
+---
+language: en
+tags:
+- tinyllama
+- language-model
+- chat
+license: apache-2.0
+---
+
+# TinyLlama 1.1B Chat Model
+
+## Model Description
+TinyLlama is a lightweight LLaMA-based model with 1.1 billion parameters, designed to perform well on conversational and text generation tasks. It has been fine-tuned specifically for chat applications, providing coherent and context-aware responses.
+
+## Training Data
+The model was trained on a diverse dataset, including web text, books, and conversational data, to make it capable of handling a wide range of language styles.
+
+## Usage
+You can use this model for conversational AI, text completion, or other natural language generation tasks. Here’s a quick example:
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("jacksonstrut/tinyllama-1.1B-chat")
+model = AutoModelForCausalLM.from_pretrained("jacksonstrut/tinyllama-1.1B-chat")
+
+input_ids = tokenizer("Hello, how are you?", return_tensors="pt").input_ids
+output = model.generate(input_ids)
+print(tokenizer.decode(output[0]))
+```
+
+## Limitations
+- The model may produce biased or inappropriate outputs as it is trained on general datasets from the internet.
+- It may not be suitable for all applications, especially those requiring factual accuracy.
+
+## License
+This model is licensed under the Apache 2.0 License.
+"""
+
+# Save the Model Card to README.md
+readme_path = os.path.join(model_dir, "README.md")
+with open(readme_path, 'w') as f:
+    f.write(model_card)
+
+# Upload the Model Card to Hugging Face Hub
+upload_file(path_or_fileobj=path_or_fileobj, path_in_repo=path_in_repo, repo_id=repo_id, token=os.getenv('HUGGINGFACE_API_TOKEN'))
+)  # Ensure the token has write permissions  # Ensure the token has write permissions
+
+# Upload files to Hugging Face Hub
+upload_file(
+    path_or_fileobj=model_weights_save_path,
+    path_in_repo="pytorch_model.bin",
+    repo_id=repo_id,
+    token=os.getenv('HUGGINGFACE_API_TOKEN')
+)
+upload_file(
+    path_or_fileobj=config_save_path,
+    path_in_repo="config.json",
+    repo_id=repo_id,
+    token=os.getenv('HUGGINGFACE_API_TOKEN')
+)
+
+# Upload the tokenizer files
+upload_file(
+    path_or_fileobj=os.path.join(model_dir, "tokenizer.model"),
+    path_in_repo="tokenizer.model",
+    repo_id=repo_id,
+    token=os.getenv('HUGGINGFACE_API_TOKEN')
+)
+upload_file(
+    path_or_fileobj=os.path.join(model_dir, "tokenizer_config.json"),
+    path_in_repo="tokenizer_config.json",
+    repo_id=repo_id,
+    token=os.getenv('HUGGINGFACE_API_TOKEN')
+)