import streamlit as st
from transformers import pipeline

# Load the translation models and the custom chatbot model
translator_ar_en = pipeline('translation_ar_to_en', model='Helsinki-NLP/opus-mt-ar-en')
translator_en_ar = pipeline('translation_en_to_ar', model='Helsinki-NLP/opus-mt-en-ar')
custom_model = pipeline('text2text-generation', model='Mohamed2002/product_model')

def translate(text, translator):
    return translator(text)[0]['translation_text']

# Streamlit app
st.title("Arabic Chatbot")

st.write("(:اهلا بحضرتك في تطبيق السمسار الالكتروني")

user_input = st.text_input("العميل:")

if st.button("Submit"):
    try:
        # Translate user input from Arabic to English
        translated_input = translate(user_input, translator_ar_en)
        
        # Generate prediction
        prediction = custom_model(translated_input)[0]['generated_text']
        
        # Translate response from English to Arabic
        translated_response = translate(prediction, translator_en_ar)
        st.write("السمسار:", translated_response)
    except Exception as e:
        st.write(f"An error occurred: {str(e)}")