import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the SmolLM model and tokenizer
# model_name = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
model_name = "HuggingFaceTB/SmolLM2-360M-Instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def smol_lm_process(job_description, system_prompt):
    # System Prompt and job description
    prompt = f"""<|im_start|>system
{system_prompt}<|im_end|>
<|im_start|>user
{job_description}<|im_end|>
<|im_start|>assistant
"""
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    output = model.generate(**inputs, max_new_tokens=512)
    response = tokenizer.decode(output[0], skip_special_tokens=False)
    # Extract the assistant's response
    start_idx = response.find("<|im_start|>assistant")
    end_idx = response.find("<|im_end|>", start_idx)
    response = response[start_idx + len("<|im_start|>assistant\n"):end_idx].strip()
    return response

def process_job_description(company_name, company_url, job_description):
    # Step 2: Extract key qualifications, skills, and requirements
    system_prompt_requirements = "Extract key qualifications, skills, and requirements from this job description. Output as bullet points. Remove benefits/salary and fluff. ONLY INCLUDE INFORMATION THAT TELLS THE USER WHAT SKILLS THE EMPLOYER SEEKS."
    role_requirements = smol_lm_process(job_description, system_prompt_requirements)

    # Step 3: Create a concise summary of the job description
    system_prompt_summary = "Create a concise 150-200 word summary of this job description. Remove company bragging and benefits information. FOCUS ON ASPECTS THAT POINT THE USER IN WHAT THE EMPLOYER WANTS FROM A CANDIDATE IN TERMS OF SKILLS, ACCOMPLISHMENTS, AND SUCH"
    clean_job_description = smol_lm_process(job_description, system_prompt_summary)

    return {
        "Company Name": company_name,
        "Company URL": company_url,
        "Original Job Description": job_description,
        "Role Requirements": role_requirements,
        "Clean Job Description": clean_job_description
    }

# Create the Gradio app
demo = gr.Blocks()

with demo:
    gr.Markdown("# Job Description Input")
    company_name = gr.Textbox(label="Company Name")
    company_url = gr.Textbox(label="Company URL")
    job_description = gr.TextArea(label="Paste Job Description")

    gr.Button("Submit").click(
        process_job_description,
        inputs=[company_name, company_url, job_description],
        outputs=gr.JSON(label="Output")
    )

if __name__ == "__main__":
    demo.launch()