import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM

# Load Hugging Face Model
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("xsanskarx/calculator-smollm2_v2")
    model = AutoModelForCausalLM.from_pretrained("xsanskarx/calculator-smollm2_v2")
    return tokenizer, model

tokenizer, model = load_model()

def calculate_expression(expression: str) -> str:
    # Encode the user input and generate the result
    input_ids = tokenizer.encode(expression, return_tensors="pt")
    output_ids = model.generate(input_ids, max_new_tokens=50)
    result = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return result

# Streamlit Interface
st.title("AI-Powered Calculator")
st.write("Enter a mathematical expression, and let the model solve it.")

expression = st.text_input("Enter Expression (e.g., 5 + 3 * (2 - 1)):")

if st.button("Calculate"):
    if expression.strip():
        try:
            result = calculate_expression(expression)
            st.success(f"Result: {result}")
        except Exception as e:
            st.error(f"Error: {str(e)}")
    else:
        st.warning("Please enter a valid expression.")