# Copyright (c) 2025 CMS Manhattan # All rights reserved. # Author: Konstantin Vladimirovich Grabko # Email: grabko@cmsmanhattan.com # Phone: +1(516)777-0945 # # This program is free software: you can redistribute it and/or modify # it under the terms of the GNU General Public License as published by # the Free Software Foundation, version 3 of the License. # # This program is distributed in the hope that it will be useful, # but WITHOUT ANY WARRANTY; without even the implied warranty of # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the # GNU General Public License for more details. # # You should have received a copy of the GNU General Public License # along with this program. If not, see . # # Additional terms: # Any commercial use or distribution of this software or derivative works # requires explicit written permission from the copyright holder. import torch import torch.nn.functional as F from transformers import GPT2TokenizerFast from gpt_pytorch import GPTPyTorch # Using the same import as in fine_tune.py import os from pathlib import Path # ============================= GENERATION SETTINGS ============================= # Temperature: Lower = more conservative and predictable answers. # Start with 0.7. Increase to 0.8 if the model starts repeating itself. TEMPERATURE = 0.7 # Top-K: Limits sampling to the K most likely tokens. # Start with 50. Increase if responses feel too boring/repetitive. TOP_K = 50 # Max Length: Maximum number of tokens to generate in one go MAX_LENGTH = 120 # ============================= PATHS ============================= # LAST_TRAINED_PATH = Path("models/gpt_last_trained.pt") LAST_TRAINED_PATH = Path("build/fine_tuning_output/epoch49/gpt_finetuned.pt") # FINAL_OUTPUT_DIR = Path("build/fine_tuning_output/final") FINAL_OUTPUT_DIR = Path("build/fine_tuning_output/epoch49/gpt_finetuned.pt") MODEL_SAVE_NAME = "gpt_finetuned.pt" device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # ============================= Chatbot CLASS ============================= class Chatbot: def __init__(self, model_path): # 1. Tokenizer print("Loading standard tokenizer (gpt2)...") self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2") self.tokenizer.pad_token = self.tokenizer.eos_token #2. Model print("Initializing model...") self.model = GPTPyTorch().to(device) self.model.eval() # Look for the latest weights: first check final folder, then last_trained load_path = None if (FINAL_OUTPUT_DIR / MODEL_SAVE_NAME).exists(): load_path = FINAL_OUTPUT_DIR / MODEL_SAVE_NAME print(f"Weights from Epoch 50 found. Loading and moving to {device}...") elif model_path.exists(): load_path = model_path print(f"Loading weights from {load_path} and moving to {device}...") if load_path: self.model.load_state_dict(torch.load(load_path, map_location=device)) else: print("Warning: No trained weights found. Using randomly initialized model.") print(f"Model successfully loaded on {device} and ready for chat!") def generate_response(self, prompt, max_length=MAX_LENGTH, temperature=TEMPERATURE, top_k=TOP_K): # Tokenize input input_ids = self.tokenizer.encode(prompt, return_tensors='pt').to(device) # Generation loop with torch.no_grad(): for _ in range(max_length): # Forward pass through the model logits, _ = self.model(input_ids) # Take logits only for the last token next_token_logits = logits[:, -1, :] # Apply temperature next_token_logits = next_token_logits / temperature # Apply Top-K sampling if top_k > 0: # Keep only the top-k most likely tokens values, indices = torch.topk(next_token_logits, top_k) # Zero out everything else (set to -inf) next_token_logits = torch.full_like(next_token_logits, float('-inf')) next_token_logits.scatter_(1, indices, values) # Convert to probabilities and sample the next token probabilities = F.softmax(next_token_logits, dim=-1) next_token = torch.multinomial(probabilities, num_samples=1) # Append generated token to the sequence input_ids = torch.cat([input_ids, next_token], dim=-1) # Stop if end-of-utterance (__eou__) or EOS token is generated generated_token = self.tokenizer.decode(next_token.squeeze().item()) if "__eou__" in generated_token or next_token.squeeze().item() == self.tokenizer.eos_token_id: break # Decode the full generated sequence output = self.tokenizer.decode(input_ids.squeeze().tolist()) # Remove the original prompt from the output response = output[len(prompt):].strip() # Clean up any leftover end-of-utterance tokens response = response.replace("__eou__", "").strip() return response def main(): # Fix for modifying globals inside the function global TEMPERATURE, TOP_K chatbot = Chatbot(LAST_TRAINED_PATH) print("\n" + "="*60) print(f"CHATBOT ACTIVATED (PPL ~2.6 / Temperature {TEMPERATURE} / Top-K {TOP_K})") print("Type 'exit' or 'quit' to quit. Use 'set temp=0.x' or 'set k=N' to change settings.") print("="*60 + "\n") while True: try: user_input = input(">>> You: ") if user_input.lower() in ['quit', 'exit']: print("Goodbye!") break # Settings commands if user_input.lower().startswith('set temp='): try: TEMPERATURE = float(user_input.split('=')[1].strip()) print(f"Temperature updated to {TEMPERATURE}") continue except ValueError: print("Invalid temperature. Use format: set temp=0.7") continue if user_input.lower().startswith('set k='): try: TOP_K = int(user_input.split('=')[1].strip()) print(f"Top-K updated to {TOP_K}") continue except ValueError: print("Invalid value. Use format: set k=50") continue print("...Generating...") response = chatbot.generate_response(user_input) print(f"Model: {response}\n") except KeyboardInterrupt: print("\nGoodbye!") break except Exception as e: print(f"An error occurred: {e}") break if __name__ == "__main__": main()