import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Step 1: Load the fine-tuned model and tokenizer
MODEL_NAME = "ridahabbash/AccR4"  # Replace with your model's Hub ID
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

# Step 2: Define the prediction function
def generate_report(prompt):
    # Tokenize input
    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)

    # Generate output
    outputs = model.generate(**inputs, max_length=128)

    # Decode and return the result
    report = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return report

# Step 3: Create the Streamlit interface
st.title("Accounting Report Generator")
st.markdown("Enter a ledger entry below, and the model will generate a report.")

# Input textbox
prompt = st.text_area("Ledger Entry", placeholder="Enter ledger details here...")

# Generate button
if st.button("Generate Report"):
    if prompt.strip() == "":
        st.error("Please enter a valid ledger entry.")
    else:
        # Generate report
        with st.spinner("Generating report..."):
            report = generate_report(prompt)
        # Display the report
        st.success("Report generated successfully!")
        st.subheader("Generated Report:")
        st.write(report)