import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

model_name = "nouamanetazi/cover-letter-t5-base"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)


def generate_cover_letter(
    name, job, company, background, experiences, max_length=300, temperature=1.0, top_p=0.9, max_time=10
):
    model_args = {
        "max_length": max_length,
        "temperature": temperature,
        "top_p": top_p,
        # "top_k": 120,
        "early_stopping": True,
        "max_time": max_time,
        "do_sample": True,  # do_sample=False to force deterministic output
        "num_return_sequences": 1,  # number of samples to return
        "min_length": 100,
        "num_beams": 4,
        # "num_beam_groups": 1,
        # "diversity_penalty": 0,
        # "repetition_penalty": 5.0,
        # "length_penalty": 0,
        # "remove_invalid_values": True,
        "no_repeat_ngram_size": 3,
    }
    # Load the tokenizer and the distilgpt2 model
    # Set up the transformers pipeline
    text_generator = pipeline(
        "text2text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1
    )
    # Generate the text
    prompt = f"coverletter name: {name} job: {job} at {company} background: {background} experiences: {experiences}"
    generated_text = text_generator(prompt, **model_args)[0]["generated_text"]
    return generated_text


title = "A Cover Letter Generator for Jobs"
description = ""
article = '<div style="text-align:center">This is a Space App for the Cover Letter</div>'
examples = None
interface = gr.Interface(
    fn=generate_cover_letter,
    inputs=[
        gr.inputs.Textbox(
            label="Your name",
            default="Sakil Ansari",
        ),
        gr.inputs.Textbox(
            label="The job you want to apply for",
            default="Data Scientist",
        ),
        gr.inputs.Textbox(
            label="The company you want to apply for",
            default="Google",
        ),
        gr.inputs.Textbox(
            lines=2,
            label="Your education/background",
            default="Master of Technology in Machine learning",
        ),
        gr.inputs.Textbox(
            lines=3,
            label="Your skills/previous experiences",
            default="I am the Author of Book and MTech in Machine Learning and achievement-driven professional with an experience of 3+ years in Data Science/Machine Learning/NLP/ Deep Learning/Data analytics. I am highly skilled in libraries like Sklearn, Numpy, Pandas, Matplotlib, Seaborn, Tensorflow, Faster-RCNN, Keras, Pytorch, FastAI, PowerBI/Tableau for Data Visualization, SQL/Oracle/NoSQL for databases and experience in NLP use cases related to named entity recognition, text summarization, text similarity, text generation.",
        ),
        gr.inputs.Slider(20, 2048, default=400, label="Max Length"),
        gr.inputs.Slider(0, 3, default=1.2, label="Temperature"),
        gr.inputs.Slider(0, 1, default=0.9, label="Top P"),
        gr.inputs.Slider(1, 200, default=20, label="Max time"),
    ],
    outputs=[gr.outputs.Textbox(type="str", label="Cover Letter")],
    title=title,
    description=description,
    examples=examples,
    article=article,
    layout="horizontal",
)
interface.launch(inline=False, debug=False)