import gradio as gr
from transformers import pipeline
from gpt4all import GPT4All

model = GPT4All("orca-mini-3b-gguf2-q4_0.gguf")

model_name = "distilbert/distilbert-base-uncased-finetuned-sst-2-english"
sentiment_analysis = pipeline("text-classification", model=model_name)

def get_sentiment(text):

    analysed_text = str(sentiment_analysis (text)[0]["label"])
    return analysed_text

def generate_prompt(user_input):
    sentiment = get_sentiment(user_input)
    if sentiment == 'POSITIVE':
        response = f"User is happy and said: {user_input}. Tell the user that it is good to know that they are happy and also write a two line poetry to celebrate their happiness."
    else:
        response = f"User is sad and said: {user_input}. Respond with a comforting message and tell them to share their issues if they would like, also write a poem to cheer them on."
    return response



def chatbot_response(input_text):
    text_prompt = generate_prompt(input_text)
    tokens = []
    with model.chat_session() as session:
        for token in model.generate(text_prompt, streaming=True):
            tokens.append(token)
    response = ''.join(tokens)
    return response

iface = gr.Interface(
    fn=chatbot_response,
    inputs=gr.components.Textbox(lines=2, placeholder="......."),
    outputs="text",
)

iface.launch(debug=True)