bsajad2/Text_generator

1import gradio as gr2import torch3import tiktoken4import os5from model_tn import GPT, GPTConfig  # Import your custom model and config6 7# --- Configuration ---8device = torch.device("cuda" if torch.cuda.is_available() else "cpu")9print(f"Using device: {device}")10 11# --- Load Model ---12model_path = "./saved_model"  # Path where your model is saved13 14# Recreate the config object used during training15config = GPTConfig(16    block_size=512,17    vocab_size=50304,18    n_layer=8,19    n_head=8,20    n_embd=768,21)22 23try:24    # Instantiate your custom GPT model25    model = GPT(config)26 27    # Load the saved state_dict28    state_dict = torch.load(os.path.join(model_path, "pytorch_model.bin"), map_location=device, weights_only=True)29    # Handle size mismatch for transformer.wpe.weight30    model_wpe_weight = model.transformer.wpe.weight31    loaded_wpe_weight = state_dict['transformer.wpe.weight']32    if model_wpe_weight.shape != loaded_wpe_weight.shape:33        print(f"Size mismatch for transformer.wpe.weight: model shape {model_wpe_weight.shape}, loaded shape {loaded_wpe_weight.shape}")34        min_size = torch.min(torch.tensor(model_wpe_weight.shape), torch.tensor(loaded_wpe_weight.shape)).tolist()35        model_wpe_weight.data[:min_size[0], :min_size[1]] = loaded_wpe_weight.data[:min_size[0], :min_size[1]]36        state_dict['transformer.wpe.weight'] = model_wpe_weight37    model.load_state_dict(state_dict)38 39    model.to(device)40    model.eval()41    print("Model loaded successfully!")42except Exception as e:43    print(f"Error loading model: {e}")44    model = None45 46# --- Load Tokenizer (tiktoken) ---47try:48    enc = tiktoken.get_encoding("gpt2")49    print("Tokenizer (tiktoken) loaded successfully!")50except Exception as e:51    print(f"Error loading tokenizer: {e}")52    enc = None53 54# --- Text Generation Function ---55def generate_text(prompt, max_length=100, temperature=0.7, top_p=0.9):56    if model is None or enc is None:57        return "Model or tokenizer failed to load. Please check the model path and files."58 59    if not prompt:60        return "Please enter a prompt."61 62    try:63        # Encode the prompt using tiktoken64        input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0).to(device)65 66        with torch.no_grad():67            for _ in range(max_length):68                # Get the logits from your custom model69                logits, _ = model(input_ids)70                logits = logits[:, -1, :] / temperature  # Apply temperature71 72                # Apply top-p sampling (nucleus sampling)73                sorted_logits, sorted_indices = torch.sort(logits, descending=True)74                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)75                sorted_indices_to_remove = cumulative_probs > top_p76                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()77                sorted_indices_to_remove[..., 0] = 078                indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)79                logits[indices_to_remove] = float('-inf')80 81                # Sample from the distribution82                probs = torch.softmax(logits, dim=-1)83                next_token = torch.multinomial(probs, num_samples=1)84 85                # Append the next token to the input_ids86                input_ids = torch.cat((input_ids, next_token), dim=1)87 88                # Check for end of sequence89                if next_token == enc.eot_token:90                    break91 92        # Decode the generated text using tiktoken93        generated_text = enc.decode(input_ids.squeeze().tolist())94        return generated_text95 96    except Exception as e:97        return f"An error occurred during text generation: {str(e)}"98 99# --- Gradio Interface ---100if model is not None and enc is not None:101    demo = gr.Interface(102        fn=generate_text,103        inputs=[104            gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="Prompt"),105            gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Max Length"),106            gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),107            gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.1, label="Top-p"),108        ],109        outputs=gr.Textbox(label="Generated Text"),110        title="Text Generator",111        description="Enter a prompt and the model will generate text.",112        examples=[113            ["I say unto you, what he hath done famously", 50, 0.7, 0.9],114            ["Your most grave belly was deliberate, Not rash like his accusers", 100, 0.8, 0.95],115        ],116    )117 118    demo.launch()119else:120    print("Model and tokenizer were not loaded successfully. Gradio interface will not be launched.")