import streamlit as st
from transformers import pipeline

# Initialize pipelines for each NLP task
ner = pipeline("ner")
qa = pipeline("question-answering")
text_gen = pipeline("text-generation")
summarization = pipeline("summarization")

def main():
  """
  This function builds the Streamlit app with user input for NLP tasks.
  """

  # Title and description for the app
  st.title("Multi-Task NLP App")
  st.write("Perform various NLP tasks on your text input.")

  # Text input field
  user_input = st.text_input("Enter Text Here:")

  # Select task from dropdown menu
  selected_task = st.selectbox("Choose NLP Task:", ["NER", "QA", "Text Generation", "Text Summarization"])

  # Perform NLP task based on selection
  if user_input and selected_task:
    if selected_task == "NER":
      analysis = ner(user_input)
      st.write("**Named Entities:**")
      for entity in analysis:
          st.write(f"- {entity['word']} ({entity['entity_group']})")
    elif selected_task == "QA":
      # Provide context (optional) for QA
      context = st.text_input("Enter Context (Optional):", "")
      if context:
          analysis = qa(question="Your question", context=context, padding="max_length")
      else:
          analysis = qa(question="Your question", context=user_input, padding="max_length")
      st.write("**Answer:**", analysis['answer'])
    elif selected_task == "Text Generation":
      # Choose generation task from another dropdown
      generation_task = st.selectbox("Choose Generation Task:", ["Text summarization (short)", "Poem", "Code"])
      if generation_task == "Text summarization (short)":
          analysis = summarization(user_input, max_length=50, truncation=True)
      else:
          # Experiment with different prompts and max_length for creative text generation
          prompt = st.text_input("Enter Prompt (Optional):", "")
          analysis = text_gen(prompt if prompt else user_input, max_length=50, truncation=True)
      st.write("**Generated Text:**", analysis[0]['generated_text'])
    else:
      analysis = summarization(user_input, max_length=100, truncation=True)
      st.write("**Summary:**", analysis[0]['summary_text'])

if __name__ == "__main__":
  main()