Update app.py

app.py

@@ -1,35 +1,30 @@
 import os
-import gradio as gr
-from huggingface_hub import login
+import streamlit as st
 from transformers import pipeline
+from dotenv import load_dotenv
 
+load_dotenv()
 # Load the gated model
 model_name = "RickyDeSkywalker/TheoremLlama"
-HF_TOKEN = os.environ.get("HF_TOKEN")
-login(HF_TOKEN)
+token = os.environ.get("HF_TOKEN")
 
-generator = pipeline('text-generation', model=model_name, token=HF_TOKEN)
+generator = pipeline('text-generation', model=model_name, token=token)
 
-# Function for generating Lean 4 code
-def generate_lean4_code(prompt):
-    result = generator(prompt, max_length=1000, num_return_sequences=1)
-    return result[0]['generated_text']
+# Streamlit app
+st.title("Lean 4 Code Generation with TheoremLlama")
+st.write("Enter a natural language prompt to generate Lean 4 code.")
 
-# Gradio Interface
-title = "Lean 4 Code Generation with TheoremLlama"
-description = "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.")
 
-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"),
-    title=title,
-    description=description
-)
+# 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')
 
-# Launch the Gradio app
-interface.launch()
+st.markdown("Developed by [Your Name]")