import gradio as gr
from transformers import pipeline
from transformers import AutoModelForSeq2SeqLM
from transformers import AutoTokenizer
from transformers import GenerationConfig
import re
model_name='google/flan-t5-base'
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
Examples_to_teach_model="""
Text: I hate apples
Sentiment analysis:
Sentiments: Negative
PPrint Key words: hate, aples
Text: I enjoy watching movies
Sentiment analysis:
Sentiments: Positive
PPrint Key words: enjoy, movies
Text: I'm tired of this long process
Sentiment analysis:
Sentiments: Negative
PPrint Key words: tired, long process
"""
def make_prompt(sentence):
    prompt = Examples_to_teach_model+ "Text: " + sentence + "Sentiment analysis:"
    return prompt
def split_conj(text):
    return re.sub('(but|yet|although|however|nevertheless|on the other hand|still|though)', "|", text).split('|')

def get_sentiment_from_llm(review_text):
    sentences = review_text.lower().split(".")
    segments=[]
    for sen in sentences:
        segments=segments+split_conj(sen)
    ls_outputs=[]

    segments= [x for x in segments if len(x)>=5]
    print(segments)
    
    for seg in segments:
        input = make_prompt(seg)
        inputs = tokenizer(input, return_tensors='pt')
        output = tokenizer.decode(
        model.generate(
            inputs["input_ids"],
            max_new_tokens=100,
        )[0],
          skip_special_tokens=True)
        ls_outputs.append("\n".join(output.split('PPrint ')))
    return "\n".join(ls_outputs)
demo = gr.Blocks()
sentiment_extr = gr.Interface(
    fn=get_sentiment_from_llm,
    inputs=gr.Textbox(label="Text input", type="text"),
    outputs=gr.Textbox(label="Sentiments", type="text"),
    title="Sentiments analysis",
    description="Sentiment analysis and keywords extraction. Powered by prompt tuned flan-t5 from Google. <br> The model is run on small CPU. Please allow 2-3 minutes for longer inputs.",
)
with demo:
    gr.TabbedInterface([sentiment_extr], ["Sentiment text analysis"])
demo.launch()