import streamlit as st
from funs import *
from transformers import pipeline
import torch
from PIL import Image

feel = pipeline("text-classification", model="SamLowe/roberta-base-go_emotions") #text classifier, it feels
knower = pipeline("text-generation", model="bigscience/bloom") #text generation, it handles

st.title("Comp-anion")
st.subheader("Comp-anion is a computer companion! Upload either text or a photo from your journal to get some insight and compassion from your comp-anion!")

u_file = st.file_uploader("Choose a file")

if u_file is not None:#when file gets uploaded
    seer = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning") #image to text , it sees
    u_file = preprocess(u_file)
    if precheck(get_context(seer,u_file)):#precheck preprocessed img
        st.write("Submission:")
        the_context=get_content(seer, u_file)
        st.write(the_context) #write out the img to text 
        emotion_found = emotion(feel,the_context)
        st.write(emotion_found)
        handle(knower,emotion_found,the_context)

st.subheader("Pictures aren't your style? Paste your text below and hit analyze!")
text_box=st.text_input("Paste Your Text Here :)", value="I've had a really nice day today")

if st.button("Analyze"):
    prompt="Give advice based on these inputs emotion={} , text given={} " 
    emotion_found=emotion(feel,text_box)
    prompt = prompt.format(emotion_found , prompt)
    st.write(prompt)
    st.write(emotion_found)
    generation=handle(knower,prompt)
    st.write(generation[:len(text_box)])