app.py

import os
import pandas as pd
import streamlit as st
import time
import random
import huggingface_hub as hf
import datasets
from datasets import load_dataset
from huggingface_hub import login
import openai


# File Path
DATA_PATH = "Dr-En-space-test.csv"
DATA_REPO = "M-A-D/dar-en-space-test"

st.set_page_config(layout="wide")

api = hf.HfApi()
access_token_write = "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"
login(token=access_token_write)

# Load data
def load_data():
    return pd.DataFrame(load_dataset(DATA_REPO,download_mode="force_redownload",split='test'))

def save_data(data):
    data.to_csv(DATA_PATH, index=False)
    # to_save = datasets.Dataset.from_pandas(data)
    api.upload_file(
    path_or_fileobj="./Dr-En-space-test.csv",
    path_in_repo="Dr-En-space-test.csv",
    repo_id=DATA_REPO,
    repo_type="dataset",
)
    # to_save.push_to_hub(DATA_REPO)

def skip_correction():
    noncorrected_sentences = st.session_state.data[(st.session_state.data.translated == True) & (st.session_state.data.corrected == False)]['sentence'].tolist()
    if noncorrected_sentences:
        st.session_state.orig_sentence = random.choice(noncorrected_sentences)
        st.session_state.orig_translation = st.session_state.data[st.session_state.data.sentence == st.session_state.orig_sentence]['translation']
    else:
        st.session_state.orig_sentence = "No more sentences to be corrected"
        st.session_state.orig_translation = "No more sentences to be corrected"

st.title("Darija Translation Corpus Collection")

if "data" not in st.session_state:
    st.session_state.data = load_data()

if "sentence" not in st.session_state:
    untranslated_sentences = st.session_state.data[st.session_state.data['translated'] == False]['sentence'].tolist()
    if untranslated_sentences:
        st.session_state.sentence = random.choice(untranslated_sentences)
    else:
        st.session_state.sentence = "No more sentences to translate"

if "orig_translation" not in st.session_state:
    noncorrected_sentences = st.session_state.data[(st.session_state.data.translated == True) & (st.session_state.data.corrected == False)]['sentence'].tolist()
    noncorrected_translations = st.session_state.data[(st.session_state.data.translated == True) & (st.session_state.data.corrected == False)]['translation'].tolist()
    
    if noncorrected_sentences:
        st.session_state.orig_sentence = random.choice(noncorrected_sentences)
        st.session_state.orig_translation = st.session_state.data.loc[st.session_state.data.sentence == st.session_state.orig_sentence]['translation'].values[0]
    else:
        st.session_state.orig_sentence = "No more sentences to be corrected"
        st.session_state.orig_translation = "No more sentences to be corrected"

if "user_translation" not in st.session_state:
    st.session_state.user_translation = ""


with st.sidebar:
    st.subheader("About")
    st.markdown("""This is app is designed to collect Darija translation corpus.""")

# tab1, tab2 = st.tabs(["Translation", "Correction"])
tab1, tab2, tab3 = st.tabs(["Translation", "Correction", "Auto-Translate"])

with tab1:
    with st.container():
        st.subheader("Original Text:")
        
        st.write('<div style="height: 150px; overflow: auto; border: 2px solid #ddd; padding: 10px; border-radius: 5px;">{}</div>'.format(st.session_state.sentence), unsafe_allow_html=True)


    st.subheader("Translation:")
    st.session_state.user_translation = st.text_area("Enter your translation here:", value=st.session_state.user_translation)
    
    if st.button("💾 Save"):
        if st.session_state.user_translation:
            st.session_state.data.loc[st.session_state.data['sentence'] == st.session_state.sentence, 'translation'] = st.session_state.user_translation
            st.session_state.data.loc[st.session_state.data['sentence'] == st.session_state.sentence, 'translated'] = True
            save_data(st.session_state.data)
            
            st.session_state.user_translation = ""  # Reset the input value after saving
            
            # st.toast("Saved!", icon="👏")
            st.success("Saved!")
            
            # Update the sentence for the next iteration.
            untranslated_sentences = st.session_state.data[st.session_state.data['translated'] == False]['sentence'].tolist()
            if untranslated_sentences:
                st.session_state.sentence = random.choice(untranslated_sentences)
                
            else:
                st.session_state.sentence = "No more sentences to translate"
            
            time.sleep(0.5)
            # Rerun the app 
            st.rerun()

with tab2:
    with st.container():
        st.subheader("Original Darija Text:")
        st.write('<div style="height: 150px; overflow: auto; border: 2px solid #ddd; padding: 10px; border-radius: 5px;">{}</div>'.format(st.session_state.orig_sentence), unsafe_allow_html=True)

    with st.container():
        st.subheader("Original English Translation:")
        st.write('<div style="height: 150px; overflow: auto; border: 2px solid #ddd; padding: 10px; border-radius: 5px;">{}</div>'.format(st.session_state.orig_translation), unsafe_allow_html=True)
    
    st.subheader("Corrected Darija Translation:")
    corrected_translation = st.text_area("Enter the corrected Darija translation here:")

    if st.button("💾 Save Translation"):
        if corrected_translation:
            st.session_state.data.loc[st.session_state.data['sentence'] == st.session_state.orig_sentence, 'translation'] = corrected_translation
            st.session_state.data.loc[st.session_state.data['sentence'] == st.session_state.orig_sentence, 'correction'] = corrected_translation
            st.session_state.data.loc[st.session_state.data['sentence'] == st.session_state.orig_sentence, 'corrected'] = True
            save_data(st.session_state.data)

            st.success("Saved!")

            # Update the sentence for the next iteration.
            noncorrected_sentences = st.session_state.data[(st.session_state.data.translated == True) & (st.session_state.data.corrected == False)]['sentence'].tolist()
            # noncorrected_sentences = st.session_state.data[st.session_state.data['corrected'] == False]['sentence'].tolist()
            if noncorrected_sentences:
                st.session_state.orig_sentence = random.choice(noncorrected_sentences)
                st.session_state.orig_translation = st.session_state.data[st.session_state.data.sentence == st.session_state.orig_sentence]['translation']

            else:
                st.session_state.orig_translation = "No more sentences to be corrected"

            corrected_translation = ""  # Reset the input value after saving

    st.button("⏩ Skip to the Next Pair", key="skip_button", on_click=skip_correction)

with tab3:
    st.subheader("Auto-Translate")

    # User input for OpenAI API key
    openai_api_key = st.text_input("Paste your OpenAI API key:")
   
    if st.button("Auto-Translate 10 Samples"):
        if openai_api_key:
            openai.api_key = openai_api_key

            # Get 10 samples from the dataset for translation
            samples_to_translate = st.session_state.data.sample(10)['sentence'].tolist()

            # System prompt for translation assistant
            translation_prompt = """
            You are a helpful AI-powered translation assistant designed for users seeking reliable translation assistance. Your primary function is to provide context-aware translations from Moroccan Arabic (Darija) to English.
            """

            auto_translations = []

            for sentence in samples_to_translate:
                # Create messages for the chat model
                messages = [
                    {"role": "system", "content": translation_prompt},
                    {"role": "user", "content": f"Translate the following sentence to English: '{sentence}'"}
                ]

                # Perform automatic translation using OpenAI GPT-3.5-turbo model
                response = openai.ChatCompletion.create(
                    model="gpt-3.5-turbo",
                    messages=messages,
                    api_key=openai_api_key
                )

                # Extract the translated text from the response
                translated_text = response.choices[0].message['content'].strip()

                # Append the translated text to the list
                auto_translations.append(translated_text)

            # Update the dataset with auto-translations
            st.session_state.data.loc[
                st.session_state.data['sentence'].isin(samples_to_translate),
                'translation'
            ] = auto_translations

            # Save the updated dataset
            save_data(st.session_state.data)

            st.success("Auto-Translations saved!")

        else:
            st.warning("Please paste your OpenAI API key.")