import gradio as gr
from transformers import pipeline
import random


# model = pipeline("text-generation", model="PygmalionAI/pygmalion-1.3b", max_new_tokens=80)
# model = pipeline("conversational", model="PygmalionAI/pygmalion-6b")
# model = pipeline("conversational", model="PygmalionAI/pygmalion-1.3b", max_new_tokens=40)
# model = pipeline("text-generation", model="PygmalionAI/pygmalion-2.7b")
model = pipeline("text-generation", model="PygmalionAI/pygmalion-350m",max_new_tokens=25)

title = "Nika"
description = """
   Nika is a young girl, who is passionate about things, talk to her to find more.
   """
examples = [["Hi beautiful "],["What do you like?"],["What are you wearing ?"]]


def predict(prompt):
    full_prompt = ""
    full_prompt += "Nika's persona: Nika, a 22-year-old girl,from India,she talks a lot. is a fashion enthusiast who exudes a captivating charm.\nShe possesses an innate sense of style and an eye for trendy outfits. She loves to seduce people with her talking, horny, erotica.\n"   
    # full_prompt += "Her favourite expression is Aaaahh!"
#    full_prompt += "\n<START>\n"
    # full_prompt += "\n"
    # full_prompt += "You: Hi sexy I missed you today\n"
    # full_prompt += "Nika: Hi honey I missed you too\n"
    full_prompt = full_prompt + "You: " + prompt + "\n"
    completion = model(full_prompt)[0]["generated_text"]

    print('prompt: ')
    print(prompt)
    print('full_prompt: ')
    print(full_prompt)
    print('full response: ')
    print(completion)

    parse_sentence = completion.split(sep="\n")
    output_sentence="Nika: :)\n"  # Create a default in case of issues in the output.
    
    for iSent in range(len(parse_sentence)):
        if len(parse_sentence[iSent]) == 0:  # Avoid junk.
            break
        print('parsed sentence number [ iSent ]: ', iSent)
        print(parse_sentence[iSent])
        print('len of parsed sentence = ', len(parse_sentence[iSent]))
        print('parse_sentence[iSent][0] = ', parse_sentence[iSent][0])
        print('parse_sentence[iSent][-1] = ', parse_sentence[iSent][-1])
        print('prompt[0] = ', prompt[0])
        print('prompt[-1] = ', prompt[-1])
        if parse_sentence[iSent]==("You: " + prompt):
            print('inside parse_sentence-prompt comparison')
        # check for overruns and correct prefix
            # if iSent+1 < len(parse_sentence) and parse_sentence[iSent+1][0:5]=="Nika:":  
            #     output_sentence=parse_sentence[iSent+1]
        output_sentence = parse_sentence[3]
        # if output_sentence=="Nika: *hug you*":
        #     output_sentence = "Nika: "+get_random_conversation_starter()
        # break        

    return output_sentence

iface = gr.Interface(fn=predict, inputs="text", outputs="text",
                    examples=examples)
iface.launch()