app.py

import streamlit as st
import base64
import torch
from transformers import AutoTokenizer, AutoModel
import requests
from huggingface_hub import login

# Get Hugging Face token from environment variable
hf_token = os.getenv("HF_API_KEY")
if hf_token is None:
    st.error("API Key not found. Please set the HF_API_KEY secret in your Space.")

# Set up Hugging Face login using the token from the environment
login(hf_token)
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabertv02")
model = AutoModel.from_pretrained("aubmindlab/bert-base-arabertv02")

# Helper functions
def get_base64(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()


def set_background(png_file):
    try:
        bin_str = get_base64(png_file)
        page_bg_img = f'''
        <style>
        .stApp {{
        background-image: url("data:image/png;base64,{bin_str}");
        background-size: contain;
        background-position: center;
        background-attachment: fixed;
        background-repeat: no-repeat;
        filter: brightness(0.9) saturate(1.2);
        }}
        .stApp::before {{
        content: "";
        position: absolute;
        top: 0;
        left: 0;
        right: 0;
        bottom: 0;
        background: url("data:image/png;base64,{bin_str}") center center / cover;
        filter: brightness(0.6);
        z-index: -1;
        }}
        </style>
        '''
        st.markdown(page_bg_img, unsafe_allow_html=True)
    except FileNotFoundError:
        st.error("Background image file not found.")

st.set_page_config(page_title="رحلات الزمان والمكان مع علّام", page_icon=":books:", layout="wide")
set_background('background_image.png')

class CustomRetriever:
    def __init__(self, tokenizer, model, max_length=512):
        self.tokenizer = tokenizer
        self.model = model
        self.max_length = max_length
        self.embeddings = None

    def tokenize_and_embed(self, text):
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=self.max_length, padding=True)
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.last_hidden_state.mean(dim=1)

    def set_corpus_embeddings(self, corpus):
        """Precomputes and stores embeddings for a corpus."""
        self.embeddings = []
        for passage in corpus:
            self.embeddings.append(self.tokenize_and_embed(passage))
        self.embeddings = torch.vstack(self.embeddings)

    def retrieve(self, question):
        """Generates query embeddings and retrieves top matching passages."""
        query_embedding = self.tokenize_and_embed(question)
        similarities = torch.nn.functional.cosine_similarity(query_embedding, self.embeddings)
        
        # Adjust top_k to handle fewer passages than k
        top_k_count = min(3, len(self.embeddings))  # Use min to avoid index out of range
        top_k = similarities.topk(top_k_count).indices
        return top_k

# IAM Token generation for IBM Cloud API
def get_iam_token(api_key):
    auth_url = "https://iam.cloud.ibm.com/identity/token"
    response = requests.post(auth_url, data={
        'apikey': api_key,
        'grant_type': 'urn:ibm:params:oauth:grant-type:apikey'
    }, headers={
        'Content-Type': 'application/x-www-form-urlencoded',
        'Accept': 'application/json'
    })
    if response.status_code == 200:
        return response.json()["access_token"]
    else:
        raise Exception(f"Error: {response.status_code} - {response.text}")

# GenerateAnswer class
def generate_answer(context, question, iam_token):
# Examples to demonstrate the structure the model should follow
    examples = """

    Q: ماهي اقسام الكلام؟

    A: كان علاء الدين شاباً يحب المغامرات والقصص، وفي أحد الأيام، أمسك بمصباحه السحري وتمنى أن يطير به إلى عالم الكلمات والحروف. فإذا به يرتفع في الهواء، حاملاً علاء الدين في رحلة عبر السماء.
    وصل علاء الدين إلى قصرٍ بديع، كان مبنيّاً من كلمات متلألئة. وعندما دخل، وجد ثلاثة أصدقاء يعيشون فيه: الأول كان اسمه "الاسم"، والثاني "الفعل"، والثالث "الحرف".
    قال "الاسم": "أنا المسؤول عن كل الأسماء التي تراها حولك، من إنسان إلى شجرة، إلى كتاب، كل شيء له اسم أعطيه إياه." أما "الفعل"، فقال: "وأنا أعطي الأشياء القدرة على الحركة والعمل، فأقول إن أحمد يقرأ، أو الطائرة تطير." وبينما كانا يتحدثان، قال "الحرف": "أنا أربط الكلمات ببعضها لتصبح جملة ذات معنى، فأنا كالغراء الذي يجمع الحروف معاً."
    جلس علاء الدين مع أصدقائه الجدد، وتحدثوا طويلاً عن الكلمات والجمل. فهم علاء الدين أن الكلمات تشبه قطع الليجو، وعندما نركبها معاً نحصل على صورة كاملة. وأن "الاسم" هو الأساس، و"الفعل" هو الحركة، و"الحرف" هو الرابط بينهما.
    عاد علاء الدين إلى أرضه وهو يحمل في قلبه الكثير من المعرفة عن الكلمات. وتذكر دائماً أن الكلمات هي أقوى سلاح لدينا، وأن علينا أن نستخدمها بحكمة وبناء.
    هل أعجبتك القصة يا صغيري؟ تذكر دائماً أن تقرأ الكثير من الكتب، لأن القراءة توسع مداركك وتزيد.

    """

    # Characters list for storytelling context
    characters_list = """
    - "شهرزاد": الحكواتية الشهيرة التي تروي القصص بمهارة لتسحر الملك شهريار، يمكنها أن تروي للأطفال قصة جديدة لتوضيح المفهوم.
    - "شهريار": الملك الذي يستمع للقصص، ويمكن أن يمثل الطفل الذي يريد معرفة المزيد ويطرح الأسئلة.
    - "السندباد البحري": البحار الذي يسافر إلى أماكن بعيدة ويواجه تحديات مختلفة، يمكنه أخذ الأطفال في مغامرات لفهم المفاهيم الصعبة.
    - "علاء الدين": الشاب الذي يملك المصباح السحري، يمكنه استخدام الجني لمساعدة الأطفال على فهم الأمور السحرية أو الغامضة.
    - "مرجانة": الفتاة الذكية التي تساعد علي بابا، يمكنها أن تقدم حلولًا سريعة وأفكارًا مفيدة.
    """

    # Instructions for how the response should be generated
    instructions = f"""
    أنت معلم للأطفال وتريد أن تجعل عملية التعلم مشوقة ومليئة بالخيال. أجب على السؤال باستخدام أسلوب سرد يشبه قصص ألف ليلة وليلة.
    اختر من الشخصيات التالية لتساعدك في الشرح:
    {characters_list}
    احرص على:
    - استخدام شخصية أو أكثر لجعل الشرح ممتعًا.
    - تقديم مغامرة بسيطة مع وصف المكان والأحداث لجعل المشهد حيويًا وشيقًا.
    - التأكد من أن الشرح مناسب للأطفال بين 6 و12 سنة باستخدام جمل قصيرة وكلمات بسيطة.
    - إشراك الطفل في القصة عبر طرح أسئلة، مثل: "هل ترى الفرق الآن؟" أو "هل ترغب في استكشاف المزيد؟".
    - إنهاء القصة بنبرة إيجابية ومشجعة.

        هذا مثال يجب أن تقتدي به عند الإجابة:
    {examples}
    """

    # Construct the prompt with the retrieved context, question, and instructions

    input_text = f"""
    السؤال: {question}
    \nالمعلومات المرجعية:\n{context}
    \nالتعليمات:\n{instructions}
    \nالإجابة:"""

    # API endpoint and model settings
    url = "https://eu-de.ml.cloud.ibm.com/ml/v1/text/generation?version=2023-05-29"
    model_id = "sdaia/allam-1-13b-instruct"
    project_id = "ed6b7fdf-5e8e-4bbd-8f93-356d126fc962"
    headers = {
        "Accept": "application/json",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {iam_token}"
    }

    body = {
        "input": input_text,
        "parameters": {
            "decoding_method": "sample",
            "max_new_tokens": 900,
            "temperature": 0.7,
            "top_k": 50,
            "top_p": 0.9,
            "repetition_penalty": 1.15
        },
        "model_id": model_id,
        "project_id": project_id
    }

    response = requests.post(url, headers=headers, json=body)

    if response.status_code == 200:
        data = response.json()
        if "results" in data and len(data["results"]) > 0:
            generated_text = data["results"][0].get("generated_text", "No generated text found")
            return generated_text
        else:
            return "No generated text found."
    else:
        return f"Error: {response.status_code} - {response.text}"

class ArabicRAG:
    def __init__(self, corpus, num_passages=3):
        # Set corpus as an attribute for easy access
        self.corpus = corpus
        self.num_passages = num_passages

        # Initialize the retriever with AraBERT
        max_length = 512
        self.retriever = CustomRetriever(tokenizer, model, max_length)
        self.retriever.set_corpus_embeddings(self.corpus)  # Precompute embeddings for the corpus

    def forward(self, question, iam_token):
        """
        Retrieve relevant context based on the user question and generate an answer.
        """
        relevant_indices = self.retriever.retrieve(question)
        context = "\n".join([self.corpus[i] for i in relevant_indices])
        prediction = generate_answer(context=context, question=question, iam_token=iam_token)
        return prediction

# API Key and IAM Token setup
api_key = 'l9cHEdqwQcXTGQ5toy6w02ogU8KR89g3w94ojrI8mgN1'
iam_token = get_iam_token(api_key)

# Split passages and initialize ArabicRAG
with open('cleaned_fixed_text.txt', 'r', encoding='utf-8') as f:
    cleaned_fixed_text = f.read()

passages = cleaned_fixed_text.split(". ")  # Splitting by double newline for demonstration
compiled_rag = ArabicRAG(corpus=passages, num_passages=3)

# Define the main page layout
def main_page():
    # CSS Styling for the button
    st.markdown('''
    <style>
    .start-button {
    display: flex;
    justify-content: center;
    align-items: center;
    position: fixed;
    bottom: 170px;
    left: 50%;
    transform: translateX(-50%);
    width: auto; max-width: 100%;
    }
    .start-button button {
    background-color: rgba(150, 98, 179, 0.8);  /* Add some transparency but not fully transparent */
    color: white;
    border: none;
    padding: 15px 100px;
    font-size: 30px;
    border-radius: 30px;
    box-shadow: inset 4px 4px 8px rgba(0, 0, 0, 0.3),
                4px 4px 8px rgba(0, 0, 0, 0.2);
    cursor: pointer;
    }
    .start-button button:hover {
        background-color: #7A5FA9;
    }
    </style>
    ''', unsafe_allow_html=True)

    # Add a button in the center of the screen with styling and make it functional
    st.markdown('<div class="start-button">', unsafe_allow_html=True)
    if st.button("ابدأ المغامرة", key="start_button_key"):
        st.session_state['page'] = 'chatbot'
    st.markdown('</div>', unsafe_allow_html=True)

# Define the chatbot page layout
def chatbot_page():
    # Title for the chatbot page
    set_background('chat_page.png')
    # CSS for fixed chat display area and input area
    st.markdown('''
        <style>
        /* Chat display area */
        .chat-display {
            position: fixed;
            top: 30px;
            right: 20px;
            width: 70%; /* Adjust width as desired */
            height: 80vh;
            padding: 15px;
            background-color: #f9f9f9;
            border: 1px solid #ddd;
            border-radius: 10px;
            overflow-y: auto;
            box-shadow: 2px 2px 10px rgba(0, 0, 0, 0.1);
        }
        /* User and bot message styles */
        .chat-bubble-user {
            background-color: #d1e7dd;
            padding: 10px;
            margin: 5px;
            border-radius: 10px;
            text-align: left;
        }
        .chat-bubble-bot {
            background-color: #f8d7da;
            color: white;
            padding: 10px;
            margin: 5px;
            border-radius: 10px;
            text-align: left;
        }
        /* Fixed input area at the right side */
        .input-container {
            position: fixed;
            bottom: 20px;
            right: 20px;
            width: 30%;
            display: flex;
            flex-direction: column;
            gap: 10px;
        }
        .chat-input {
            width: 100%;
            padding: 15px;
            font-size: 18px;
            border-radius: 10px;
            border: 1px solid #ccc;
            box-shadow: 2px 2px 8px rgba(0, 0, 0, 0.1);
        }
        .chat-button {
            padding: 10px;
            font-size: 18px;
            border-radius: 10px;
            border: none;
            cursor: pointer;
            background-color: #7A5FA9;
            color: white;
        }
        .chat-button:hover {
            background-color: #5a3f73;
        }
        </style>
    ''', unsafe_allow_html=True)

    # Initialize conversation history
    if 'conversation' not in st.session_state:
        st.session_state['conversation'] = []

    # Capture and process input
    user_question = st.text_input("اكتب رسالتك هنا...", key="user_question_input")

    # Display two buttons: one for generating the answer, and one for evaluation
    st.markdown('<div class="input-container">', unsafe_allow_html=True)
    if st.button("إرسال السؤال", key="send_button_key"):
        if user_question:
            response = compiled_rag.forward(user_question, iam_token)
            if len(response.strip()) < 20:  # Assuming a good answer should be at least 20 characters
                st.write("The generated answer was too short, retrying...")
                response = compiled_rag.forward(user_question, iam_token)

            st.session_state['conversation'].append({"role": "user", "content": user_question})
            st.session_state['conversation'].append({"role": "bot", "content": response})
            st.experimental_rerun()

    if st.button("تقييم الاستجابة", key="evaluation_button_key"):
        st.session_state['page'] = 'evaluation'

    st.markdown('</div>', unsafe_allow_html=True)

    # Display conversation history
    for message in st.session_state['conversation']:
        if message['role'] == 'user':
            st.markdown(f'''
                <script>
                document.getElementById("chat-display").innerHTML += '<div class="chat-bubble-user">{message["content"]}</div>';
                </script>
            ''', unsafe_allow_html=True)
        else:
            st.markdown(f'''
                <script>
                document.getElementById("chat-display").innerHTML += '<div class="chat-bubble-bot">{message["content"]}</div>';
                </script>
            ''', unsafe_allow_html=True)

# Define the evaluation page layout
def evaluation_page():
    set_background('human_rate.png')
    st.markdown('''
        <style>
        .star-rating {
            display: flex;
            flex-direction: row-reverse;
            justify-content: center;
            position: fixed;
            bottom: 100px;
            left: 50%;
            transform: translateX(-50%);
        }

        .star-rating input[type="radio"] {
            display: none;
        }

        .star-rating label {
            font-size: 2em;
            color: #ccc;
            cursor: pointer;
            transition: color 0.2s;
        }

        .star-rating input[type="radio"]:checked ~ label {
            color: #FFD700;
        }

        .star-rating label:hover,
        .star-rating label:hover ~ label {
            color: #FFD700;
        }
        </style>
        
        <div class="star-rating">
            <input type="radio" id="star5" name="rating" value="5"><label for="star5">★</label>
            <input type="radio" id="star4" name="rating" value="4"><label for="star4">★</label>
            <input type="radio" id="star3" name="rating" value="3"><label for="star3">★</label>
            <input type="radio" id="star2" name="rating" value="2"><label for="star2">★</label>
            <input type="radio" id="star1" name="rating" value="1"><label for="star1">★</label>
        </div>
        
        <script>
        const stars = document.querySelectorAll('.star-rating input');
        stars.forEach(star => {
            star.addEventListener('change', (event) => {
                const rating = event.target.value;
                window.parent.postMessage({type: 'rating', value: rating}, '*');
            });
        });
        </script>
    ''', unsafe_allow_html=True)

# Main loop
if 'page' not in st.session_state:
    st.session_state['page'] = 'main'

if st.session_state['page'] == 'main':
    main_page()
elif st.session_state['page'] == 'chatbot':
    chatbot_page()
elif st.session_state['page'] == 'evaluation':
    evaluation_page()