pages/Conversion.py

import streamlit as st
from utils import load_model, convert_to_torchscript, convert_to_onnx, get_hf_token

st.title("🔧 Model Conversion")

# Load the HF token from utils
hf_token = get_hf_token()

# Load the model
model_path = "fine_tuned_model.pt"
tokenizer, model = load_model("google/gemma-3-1b-it", hf_token, model_path)

conversion_option = st.selectbox("Select Conversion Format", ["TorchScript", "ONNX"])

if st.button("Convert Model"):
    if conversion_option == "TorchScript":
        with st.spinner("Converting to TorchScript..."):
            ts_model = convert_to_torchscript(model)
            st.success("Model converted to TorchScript!")
            
    elif conversion_option == "ONNX":
        with st.spinner("Converting to ONNX..."):
            onnx_path = convert_to_onnx(model)
            st.success("Model converted to ONNX!")