# Import the necessary libraries
import gradio as gr
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel
import os

# Get the Hugging Face token from the environment variable
HF_TOKEN = os.environ.get("HF_TOKEN")

# Load the tokenizer and model
tokenizer = GPT2Tokenizer.from_pretrained('gpt2', use_auth_token=HF_TOKEN)
model = GPT2LMHeadModel.from_pretrained('skylersterling/TopicGPT', use_auth_token=HF_TOKEN)
model.eval()
model.to('cpu')

# Define the function that generates text from a prompt
def generate_text(prompt, temperature, top_p):
    input_tokens = tokenizer.encode(prompt, return_tensors='pt')
    input_tokens = input_tokens.to('cpu')
    
    generated_text = prompt  # Start with the initial prompt

    for _ in range(80):  # Adjust the range to control the number of tokens generated
        with torch.no_grad():
            outputs = model(input_tokens)
            predictions = outputs.logits / temperature
            sorted_logits, sorted_indices = torch.sort(predictions[:, -1, :], descending=True)
            cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0
            indices_to_remove = sorted_indices[sorted_indices_to_remove]
            predictions[:, -1, indices_to_remove] = -float('Inf')
            next_token = torch.multinomial(torch.softmax(predictions[:, -1, :], dim=-1), 1)
        
        input_tokens = torch.cat((input_tokens, next_token), dim=1)
        
        decoded_token = tokenizer.decode(next_token.item())
        generated_text += decoded_token  # Append the new token to the generated text
        if decoded_token == "#":  # Stop if the end of sequence token is generated
            break
        yield generated_text  # Yield the entire generated text so far

# Create a Gradio interface with a text input, sliders for temperature and top_p, and a text output
interface = gr.Interface(
    fn=generate_text,
    inputs=[
        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=1.0, label="Temperature"),
        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=0.9, label="Top-p")
    ],
    outputs='text',
    live=False
)

interface.launch()