import torch
import gradio as gr
from transformers import AutoTokenizer
from model_smol2 import LlamaForCausalLM, config_model

# Instantiate the model
model = LlamaForCausalLM(config_model)

# Load the checkpoint
checkpoint_path = "/Users/shriti/Downloads/Assign13_ERAV3/deply/final_checkpoint.pt"
checkpoint = torch.load(checkpoint_path, map_location="cpu")
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Load tokenizer (replace with the appropriate tokenizer if you're using a custom one)
# Load the tokenizer
TOKENIZER_PATH = "HuggingFaceTB/cosmo2-tokenizer"
tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token if tokenizer.eos_token else "[PAD]"


# Text generation function
def generate_text(
    prompt, max_length=50, temperature=0.7, top_k=50, repetition_penalty=1.2, n_gram_block=2
):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    generated_tokens = input_ids[0].tolist()

    with torch.no_grad():
        for _ in range(max_length):
            outputs = model(input_ids)  # model outputs
            
            # Check if the output is a dictionary with logits
            if isinstance(outputs, dict) and 'logits' in outputs:
                logits = outputs['logits'][:, -1, :]
            else:
                # If not, treat the output as a plain tensor
                logits = outputs[:, -1, :]

            # Repetition penalty
            for token_id in set(generated_tokens):
                logits[:, token_id] /= repetition_penalty

            # n-gram blocking
            if len(generated_tokens) >= n_gram_block:
                n_gram = tuple(generated_tokens[-n_gram_block:])
                for token_id in set(generated_tokens):
                    if generated_tokens[-n_gram_block:] == list(n_gram):
                        logits[:, token_id] -= 1e9

            logits /= temperature
            top_k_logits, top_k_indices = torch.topk(logits, top_k, dim=-1)
            probs = torch.softmax(top_k_logits, dim=-1)

            next_token_idx = torch.multinomial(probs, num_samples=1)
            next_token = top_k_indices[0, next_token_idx[0]]

            generated_tokens.append(next_token.item())
            input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)

            if next_token.item() == tokenizer.eos_token_id:
                break

    return tokenizer.decode(generated_tokens, skip_special_tokens=True)


# Gradio UI
def generate_response(prompt, max_length, temperature, top_k, repetition_penalty, n_gram_block):
    return generate_text(prompt, max_length, temperature, top_k, repetition_penalty, n_gram_block)

with gr.Blocks() as demo:
    gr.Markdown("# Smol2 Text Generator")
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Input Prompt", placeholder="Enter your text prompt here...")
            max_length = gr.Slider(label="Max Length", minimum=10, maximum=200, value=50)
            temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.5, value=0.7, step=0.1)
            top_k = gr.Slider(label="Top K", minimum=10, maximum=100, value=50, step=1)
            repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.2, step=0.1)
            n_gram_block = gr.Slider(label="N-Gram Blocking", minimum=1, maximum=5, value=2, step=1)
            generate_button = gr.Button("Generate Text")
        with gr.Column():
            output_text = gr.Textbox(label="Generated Text", lines=10)

    generate_button.click(
        generate_response,
        inputs=[prompt_input, max_length, temperature, top_k, repetition_penalty, n_gram_block],
        outputs=[output_text],
    )

demo.launch()