import re
import streamlit as st
from transformers import pipeline

# Initialize the chat history
history = []

def clean_text(text):
    return re.sub('[^a-zA-Z\s]', '', text).strip()

# Load DistilBert model
model = pipeline("question-answering", model="distilbert-base-cased")

def generate_response(user_input):
    # Add user input to history
    history.append((user_input, ""))

    if not history:
        return ""
    
    last_user_message = history[-1][0]
    user_input = clean_text(last_user_message)

    if len(user_input) > 0:
        result = model(question=user_input, context="Placeholder text")
        answer = result['answer']
        history[-1] = (last_user_message, answer)

        return f"AI: {answer}"

st.title("Simple Chat App using DistilBert Model (HuggingFace & Streamlit)")

for i in range(len(history)):
    message = history[i][0]
    response = history[i][1]

    if i % 2 == 0:
        col1, col2 = st.beta_columns([0.8, 0.2])
        with col1:
            st.markdown(f">> {message}")
        with col2:
            st.write("")
    else:
        col1, col2 = st.beta_columns([0.8, 0.2])
        with col1:
            st.markdown(f"   {response}")
        with col2:
            st.button("Clear")

new_message = st.text_area("Type something...")
if st.button("Submit"):
    generated_response = generate_response(new_message)
    st.markdown(generated_response)