app.py

import streamlit as st
import random
import pandas as pd
import requests
from io import BytesIO
from PIL import Image
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import re

# Define maximum dimensions for the fortune image (in pixels)
MAX_SIZE = (400, 400)

# Initialize button click count in session state
if "button_count_temp" not in st.session_state:
    st.session_state.button_count_temp = 0

# Set page configuration
st.set_page_config(page_title="Fortuen Stick Enquiry", layout="wide")
st.title("Fortuen Stick Enquiry")

# Initialize session state variables
if "submitted_text" not in st.session_state:
    st.session_state.submitted_text = False
if "fortune_number" not in st.session_state:
    st.session_state.fortune_number = None
if "fortune_row" not in st.session_state:
    st.session_state.fortune_row = None
if "error_message" not in st.session_state:
    st.session_state.error_message = ""
if "cfu_explain_text" not in st.session_state:
    st.session_state.cfu_explain_text = ""

if "fortune_data" not in st.session_state:
    try:
        st.session_state.fortune_data = pd.read_csv("/home/user/app/detail.csv")
    except Exception as e:
        st.error(f"Error loading CSV: {e}")
        st.session_state.fortune_data = None

if "stick_clicked" not in st.session_state:
    st.session_state.stick_clicked = False

def load_finetuned_classifier_model(question):
    label_list = ["Geomancy", "Lost Property", "Personal Well-Being", "Future Prospect", "Traveling"]
    # Create a mapping dictionary to convert the default "LABEL_x" output.
    mapping = {f"LABEL_{i}": label for i, label in enumerate(label_list)}

    pipe = pipeline("text-classification", model="tonyhui2234/CustomModel_classifier_model_10")
    prediction = pipe(question)[0]['label']
    predicted_label = mapping.get(prediction, prediction)
    print(predicted_label)
    return predicted_label

# Define your inference function
def generate_answer(question, fortune):
    # Load the saved model and tokenizer
    tokenizer = AutoTokenizer.from_pretrained("/home/user/app/my_finetuned_model_2/")
    model = AutoModelForSeq2SeqLM.from_pretrained("/home/user/app/my_finetuned_model_2/")

    input_text = "Question: " + question + " Fortune: " + fortune
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
    outputs = model.generate(**inputs, max_length=256, num_beams=4, early_stopping=True, repetition_penalty=2.0, no_repeat_ngram_size=3)
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

def analysis(row_detail, classifiy, question):
    # Use the classifier's output (e.g. "Personal Well-Being") in the regex.
    pattern = re.compile(re.escape(classifiy) + r":\s*(.*?)(?:\.|$)", re.IGNORECASE)
    match = pattern.search(row_detail)
    if match:
        result = match.group(1)
        # If you want to generate a custom answer, you can call generate_answer()
        return generate_answer(question, result)
        # return result
    else:
        return "Heaven's secret cannot be revealed."

def check_sentence_is_english_model(question):
    pipe_english = pipeline("text-classification", model="papluca/xlm-roberta-base-language-detection")
    return pipe_english(question)[0]['label'] == 'en'

def check_sentence_is_question_model(question):
    pipe_question = pipeline("text-classification", model="shahrukhx01/question-vs-statement-classifier")
    return pipe_question(question)[0]['label'] == 'LABEL_1'

def submit_text_callback():
    question = st.session_state.get("user_sentence", "")
    # Clear any previous error message
    st.session_state.error_message = ""
    
    if not check_sentence_is_english_model(question):
        st.session_state.error_message = "Please enter in English!"
        st.session_state.button_count_temp = 0
        return

    if not check_sentence_is_question_model(question):
        st.session_state.error_message = "This is not a question. Please enter again!"
        st.session_state.button_count_temp = 0
        return

    if st.session_state.button_count_temp == 0:
        st.session_state.error_message = "Please take a moment to quietly reflect on your question in your mind, then click submit."
        st.session_state.button_count_temp = 1
        return

    st.session_state.submitted_text = True
    st.session_state.button_count_temp = 0  # Reset the counter once submission is accepted

    # Randomly generate a number from 1 to 100
    st.session_state.fortune_number = random.randint(1, 100)
    
    # Look up the row in the CSV where CNumber matches the generated fortune number.
    df = st.session_state.fortune_data
    row_detail = ''
    if df is not None:
        matching_row = df[df['CNumber'] == st.session_state.fortune_number]
        if not matching_row.empty:
            row = matching_row.iloc[0]
            row_detail = row.get("Detail", "No detail available.")
            st.session_state.fortune_row = {
                "Header": row.get("Header", "N/A"),
                "Luck": row.get("Luck", "N/A"),
                "Description": row.get("Description", "No description available."),
                "Detail": row_detail,
                "HeaderLink": row.get("link", None)
            }
        else:
            st.session_state.fortune_row = {
                "Header": "N/A",
                "Luck": "N/A",
                "Description": "No description available.",
                "Detail": "No detail available.",
                "HeaderLink": None
            }
    print(row_detail)
    classifiy = load_finetuned_classifier_model(question)
    cfu_explain = analysis(row_detail, classifiy, question)
    # Save the returned value in session state for later display
    st.session_state.cfu_explain_text = cfu_explain

def load_and_resize_image(path, max_size=MAX_SIZE):
    try:
        img = Image.open(path)
        img.thumbnail(max_size, Image.Resampling.LANCZOS)
        return img
    except Exception as e:
        st.error(f"Error loading image: {e}")
        return None

def download_and_resize_image(url, max_size=MAX_SIZE):
    try:
        response = requests.get(url)
        response.raise_for_status()
        image_bytes = BytesIO(response.content)
        img = Image.open(image_bytes)
        img.thumbnail(max_size, Image.Resampling.LANCZOS)
        return img
    except Exception as e:
        st.error(f"Error loading image from URL: {e}")
        return None

def stick_enquiry_callback():
    st.session_state.stick_clicked = True

# Main layout: Left (input) and Right (fortune display)
left_col, _, right_col = st.columns([3, 1, 5])

# ---- Left Column ----
with left_col:
    left_top = st.container()
    left_bottom = st.container()
    with left_top:
        st.text_area("Enter your question in English", key="user_sentence", height=150)
        st.button("submit", key="submit_button", on_click=submit_text_callback)
        if st.session_state.error_message:
            st.error(st.session_state.error_message)
    if st.session_state.submitted_text:
        with left_bottom:
            for _ in range(5):
                st.write("")
            col1, col2, col3 = st.columns(3)
            with col2:
                st.button("Cfu Explain", key="stick_button", on_click=stick_enquiry_callback)
            if st.session_state.stick_clicked:
                # Display the explanation text saved from analysis()
                st.text_area(' ', value=st.session_state.cfu_explain_text, height=300, disabled=True)

# ---- Right Column ----
with right_col:
    with st.container():
        col_left, col_center, col_right = st.columns([1, 2, 1])
        with col_center:
            if st.session_state.submitted_text and st.session_state.fortune_row:
                header_link = st.session_state.fortune_row.get("HeaderLink")
                if header_link:
                    img_from_url = download_and_resize_image(header_link)
                    if img_from_url:
                        st.image(img_from_url, use_container_width=False)
                    else:
                        img = load_and_resize_image("/home/user/app/error.png")
                        if img:
                            st.image(img, use_container_width=False)
                else:
                    img = load_and_resize_image("/home/user/app/error.png")
                    if img:
                        st.image(img, use_container_width=False)
            else:
                img = load_and_resize_image("/home/user/app/fortune.png")
                if img:
                    st.image(img, caption="Your Fortune", use_container_width=False)
    with st.container():
        if st.session_state.fortune_row:
            luck_text = st.session_state.fortune_row.get("Luck", "N/A")
            description_text = st.session_state.fortune_row.get("Description", "No description available.")
            detail_text = st.session_state.fortune_row.get("Detail", "No detail available.")
            
            summary = f"""
            <div style="font-size: 28px; font-weight: bold;">
                Fortune stick number: {st.session_state.fortune_number}<br>
                Luck: {luck_text}
            </div>
            """
            st.markdown(summary, unsafe_allow_html=True)
            
            st.text_area("Description", value=description_text, height=150, disabled=True)
            st.text_area("Detail", value=detail_text, height=150, disabled=True)