app.py

from transformers import AutoTokenizer, AutoModelForQuestionAnswering
import streamlit as st
import torch
import re

# Load model and tokenizer
model_path = "./"  # Adjust if needed
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForQuestionAnswering.from_pretrained(model_path)

# Streamlit UI
st.title("Tax QA Model")

question = st.text_area("question", "What is the tax rate?")
context = st.text_area("context", "In 2025, the income tax rate for individuals earning less than $50,000 is 10%.")

if st.button("Submit"):
    inputs = tokenizer(question, context, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    start = torch.argmax(outputs.start_logits)
    end = torch.argmax(outputs.end_logits) + 1
    answer = tokenizer.convert_tokens_to_string(tokenizer.convert_ids_to_tokens(inputs['input_ids'][0][start:end]))

    # Extract percentage values only if the question is about rates
    percentage_match = re.findall(r'\d+%', answer)
    final_answer = percentage_match[0] if percentage_match else answer

    st.text_area("output", final_answer)

if st.button("Clear"):
    st.experimental_rerun()