import os
import streamlit as st
from transformers import pipeline

# Load the gated model
model_name = "internlm/internlm2-math-plus-7b"
token = os.environ.get("HF_TOKEN")

generator = pipeline('text-generation', model=model_name, token=token, trust_remote_code=True)

# Streamlit app
st.title("Lean 4 Code Generation with TheoremLlama")
st.write("Enter a natural language prompt to generate Lean 4 code.")

# Text input for user prompt
prompt = st.text_area("Prompt", value="Prove that the sum of two even numbers is even.")

# Generate button
if st.button("Generate Lean 4 Code"):
    with st.spinner("Generating Lean 4 code..."):
        # Generate code based on the prompt
        result = generator(prompt, max_length=1000, num_return_sequences=1)
        generated_code = result[0]['generated_text']
    
    # Display the generated code
    st.code(generated_code, language='lean')

st.markdown("Developed by [Your Name]")