import os
import gradio as gr
from huggingface_hub import login
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load the gated model
#model_name = "RickyDeSkywalker/TheoremLlama"
#model_name = "internlm/internlm2-math-plus-7b"
#model_name = "deepseek-ai/DeepSeek-Prover-V1.5-RL"
model_name = "unsloth/Llama-3.2-1B-Instruct"

HF_TOKEN = os.environ.get("HF_TOKEN")
#login(HF_TOKEN)

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16, device_map="auto").eval()

# Function for generating Lean 4 code
@torch.inference_mode()
def generate_lean4_code(prompt: str):
    terminators = [tokenizer.eos_token_id, 
               tokenizer.convert_tokens_to_ids("<|eot_id|>"), 
               tokenizer.convert_tokens_to_ids("<|reserved_special_token_26|>")]
    chat = [
        {"role": "system", "content": "You are a Lean4 expert who can write good Lean4 code based on natural language mathematical theorem and proof"},
        {"role": "user", "content": prompt},
    ]
    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(model.device)
    results = model.generate(input_ids, max_new_tokens=1024, eos_token_id=terminators, do_sample=True, temperature=0.85, top_p=0.9)
    result_str = tokenizer.decode(results[0], skip_special_tokens=True)
    return result_str

# Gradio Interface
title = "Lean 4 Code Generation with TheoremLlama"
description = "Enter a natural language prompt to generate Lean 4 code."

interface = gr.Interface(
    fn=generate_lean4_code,
    inputs=gr.Textbox(
        label="Prompt",
        placeholder="Prove that the sum of two even numbers is even.",
        lines=4
    ),
    outputs=gr.Code(label="Generated Lean 4 Code"), # language="lean" unsupported in Gradio...
    title=title,
    description=description
)

# Launch the Gradio app
interface.queue().launch(ssr_mode=False)