import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("WhiteRabbitNeo/WhiteRabbitNeo-13B-v1", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained("WhiteRabbitNeo/WhiteRabbitNeo-13B-v1", trust_remote_code=True)

# Define a function to generate text with a system message
def generate_response(prompt, system_message, token):
    # Combine system message and user prompt
    full_prompt = f"{system_message}\n\n{prompt}"
    
    # Tokenize input
    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True)
    
    # Generate a response
    outputs = model.generate(
        inputs["input_ids"],
        max_length=300,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        pad_token_id=tokenizer.eos_token_id  # To prevent padding issues
    )
    
    # Decode the output text
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # If a token is provided, add it to the response
    if token:
        response += f"\n\nToken Used: {token}"
    
    return response

# Create a Gradio interface
interface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(lines=5, label="Prompt", placeholder="Type your prompt here..."),
        gr.Textbox(lines=1, label="System Message", placeholder="System message (optional)"),
        gr.Textbox(lines=1, label="Token", placeholder="Enter token (optional)"),
    ],
    outputs=gr.Textbox(label="Generated Response"),
    title="WhiteRabbitNeo Enhanced Model",
    description=(
        "This app uses the WhiteRabbitNeo-13B-v1 model to generate text responses. "
        "You can provide a system message, a prompt, and optionally include a token for custom usage."
    ),
)

# Launch the app
if __name__ == "__main__":
    interface.launch()