Committing all changes before LFS migration

chat_cli.py

@@ -0,0 +1,39 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+def chat_with_model(model_path: str):
+    # Ensure CUDA is available and set the device to use the first GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    # Load the model and tokenizer
+    model = AutoModelForCausalLM.from_pretrained(model_path).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+    # Wrap the model with DataParallel to use multiple GPUs
+    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
+        print(f"Using {torch.cuda.device_count()} GPUs!")
+        model = torch.nn.DataParallel(model)
+
+    print("You're now chatting with the model. Type 'quit' to exit.")
+
+    while True:
+        # Get user input
+        input_text = input("You: ")
+        if input_text.lower() == 'quit':
+            break
+
+        # Encode the input text
+        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+        
+        # Generate a response
+        with torch.no_grad():
+            generated_text_samples = model.generate(input_ids, max_length=50, pad_token_id=tokenizer.eos_token_id)
+
+        # Decode and print the model's response
+        response_text = tokenizer.decode(generated_text_samples[0], skip_special_tokens=True)
+        print("AI:", response_text)
+
+if __name__ == "__main__":
+    model_path = '/home/energyxadmin/UI2/merge'
+    chat_with_model(model_path)

merge_script.py

@@ -0,0 +1,63 @@
+import torch
+from peft import PeftModel  # Ensure you have 'peft' library or modify according to your setup
+import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
+import argparse
+from utils import get_logger  # Ensure this is implemented in your environment
+import json
+
+logger = get_logger("merge", "info")
+
+def smart_tokenizer_and_embedding_resize(tokenizer, model, custom_tokens_path=None):
+    """Resize tokenizer and embedding to accommodate new tokens."""
+    special_tokens_dict = {
+        "pad_token": "[PAD]",
+        "eos_token": "</s>",
+        "bos_token": "<s>",
+        "unk_token": "<unk>"
+    }
+    
+    # Load custom tokens if specified
+    custom_tokens = []
+    if custom_tokens_path is not None:
+        with open(custom_tokens_path, 'r') as file:
+            custom_tokens = [line.strip() for line in file.readlines()]
+
+    num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)
+    if custom_tokens:
+        num_added_toks += tokenizer.add_tokens(custom_tokens, special_tokens=True)
+    
+    model.resize_token_embeddings(len(tokenizer))
+    logger.info(f"Resized tokenizer and model embeddings. Added {num_added_toks} tokens.")
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-bm", "--base_model", type=str, default="meta-llama/Llama-2-7b-chat-hf", help="Base model name or path")
+    parser.add_argument("-lm", "--lora_model", type=str, required=True, help="Path to the Lora model directory")
+    parser.add_argument("-o", "--output", type=str, required=True, help="Output directory for the merged model")
+    parser.add_argument("--custom_tokens", type=str, default=None, help="Path to a file containing custom tokens")
+    args = parser.parse_args()
+
+    if not os.path.exists(args.lora_model):
+        raise FileNotFoundError(f"LoRA model directory {args.lora_model} not found.")
+    
+    os.makedirs(args.output, exist_ok=True)
+    
+    # Load the base model and tokenizer
+    model = AutoModelForCausalLM.from_pretrained(args.base_model)
+    tokenizer = AutoTokenizer.from_pretrained(args.base_model)
+
+    # Adjust tokenizer and model for any additional tokens
+    smart_tokenizer_and_embedding_resize(tokenizer, model, args.custom_tokens)
+
+    # Load and merge the LoRA model
+    logger.info("Loading and merging the LoRA model...")
+    lora_model = PeftModel.from_pretrained(model, args.lora_model, merge_with_base=True)
+    
+    # Save the merged model and tokenizer
+    lora_model.save_pretrained(args.output)
+    tokenizer.save_pretrained(args.output)
+    logger.info(f"Merged model saved to {args.output}")
+
+if __name__ == "__main__":
+    main()

test.py

@@ -0,0 +1,12 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_path = '/home/energyxadmin/UI2/merge'
+
+model = AutoModelForCausalLM.from_pretrained(model_path)
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# Example text generation
+input_ids = tokenizer.encode("What song did Eric Pask write or was a part of", return_tensors="pt")
+generated_text_samples = model.generate(input_ids, max_length=1000)
+
+print("Generated text:", tokenizer.decode(generated_text_samples[0], skip_special_tokens=True))

utils.py

@@ -0,0 +1,17 @@
+import logging
+from typing_extensions import Literal
+from rich.logging import RichHandler
+
+
+def get_logger(name: str, level: Literal["info", "warning", "debug"]) -> logging.Logger:
+    rich_handler = RichHandler(level=logging.INFO, rich_tracebacks=True, markup=True)
+
+    logger = logging.getLogger(name)
+    logger.setLevel(logging._nameToLevel[level.upper()])
+
+    if not logger.handlers:
+        logger.addHandler(rich_handler)
+
+    logger.propagate = False
+
+    return logger