import streamlit as st
from transformers import pipeline
import sympy as sp

# Cache the model so it's loaded only once
@st.cache_resource
def load_model():
    # Load an open-source Hugging Face model for natural language processing
    return pipeline('text-classification', model='mrm8488/t5-base-finetuned-summarize-news')

# Initialize the model
nlp_model = load_model()

# Function to check if the question is mathematical
def is_math_question(question):
    try:
        parsed_expr = sp.sympify(question)
        return True
    except (sp.SympifyError, SyntaxError):
        return False

# Function to solve mathematical questions using SymPy
def solve_math_question(question):
    try:
        # Parse and solve the mathematical expression
        solution = sp.solve(sp.sympify(question))
        return f"The solution is: {solution}"
    except Exception as e:
        return f"Error solving the equation: {e}"

# Streamlit UI
st.title("Math Chatbot (Open Source)")
st.write("Ask any mathematical question and get an answer. Non-mathematical questions will be restricted.")

# User input
question = st.text_input("Enter your mathematical question:")

# Processing the input
if st.button("Submit"):
    if is_math_question(question):
        # Solve the mathematical question
        answer = solve_math_question(question)
        st.write(f"Answer: {answer}")
    else:
        # Filter non-mathematical questions using NLP model
        nlp_result = nlp_model(question)[0]
        if nlp_result['label'] == 'Math':
            st.write("Answer: Processing your math question...")
        else:
            st.write("This chatbot only answers questions related to mathematics. Please ask a mathematical question.")