Create app.py

app.py

@@ -0,0 +1,169 @@
+import torch
+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 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')
+)