import gradio as gr

import os
import torch
from dotenv import load_dotenv
from datasets import load_dataset
from peft import AutoPeftModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM

load_dotenv()

def format_instruction(report):
    return """### Instruction: 
Classify the student into Placed/NotPlaced based on his/her college report details. The report includes marks scored by the student in various courses and extra curricular activities taken by them.

### Report:
{report}

### Label:
"""

def postprocess(outputs, tokenizer, prompt):
    outputs = outputs.numpy()
    outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    output = outputs[0][len(prompt):]
    
    return output


def run_model(report):
    # load dataset and select a random sample
    prompt = format_instruction(report)

    # load base LLM model, LoRA params and tokenizer
    model = AutoPeftModelForCausalLM.from_pretrained(
        os.getenv("Model_Repo_ID"),
        low_cpu_mem_usage=True,
        torch_dtype=torch.float16,
        #load_in_4bit=True
    )
    tokenizer = AutoTokenizer.from_pretrained(os.getenv("Model_Repo_ID"))
    input_ids = tokenizer(prompt, return_tensors="pt", truncation=True).input_ids.cpu()
    
    # inference
    with torch.inference_mode():
        outputs = model.generate(
            input_ids=input_ids, 
            max_new_tokens=800, 
            do_sample=True, 
            top_p=0.9,
            temperature=0.9
        )

    return postprocess(outputs, tokenizer, report)


iface = gr.Interface(
        fn=run_model, 
        inputs="text", 
        outputs=[gr.Textbox(label="Status")],
        examples=[
            ["Meet John Doe, a college student who has demonstrated impressive academic achievement and valuable work experience. John has a cumulative GPA of 6.5 in his university, indicating a strong academic record. Additionally, John has completed one internship and three projects, demonstrating his practical skills and ability to apply his knowledge in real-world settings. John also possesses a certification in workshops, further enhancing his skills and knowledge. His aptitude test score of 68 suggests a strong cognitive ability, while his soft skills rating of 4.0 indicates a strong interpersonal and communication ability. John is also involved in extracurricular activities, indicating a well-rounded personality and a desire to give back to the community. Finally, John has undergone placement training, suggesting a strong interest in pursuing a professional career. Overall, John Doe presents a strong and well-rounded profile, with a combination of academic achievement, practical experience, and valuable skills and certifications."],
            ["John Smith is a college student with a strong academic background. He has an impressive CGPA of 6.5, which is above average in most universities. John has also completed one internship and one project, which demonstrates his practical experience and ability to apply theoretical concepts in real-life scenarios. In addition to his academic and practical accomplishments, John has also completed one workshop or certification, which shows his willingness to continue learning and improving his skills. John performed well in his aptitude test, scoring a respectable 70. His soft skills rating of 3.5 is also relatively high, indicating that he is well-versed in communication, teamwork, and other important interpersonal skills. It is worth noting that John does not participate in extracurricular activities, which may limit his overall exposure to diverse perspectives and experiences. However, he has undergone placement training, which could help him develop important skills such as interviewing, networking, and job search strategies. John's academic achievements in SSC and HSC are also remarkable, with SSC marks of 60 and HSC marks of 75. Overall, John Smith appears to be a well-rounded and talented student with a strong academic record and practical experience. His high CGPA, internships, projects, and workshops or certifications are particularly impressive, and his aptitude test score and soft skills rating are also noteworthy."],
            ["Dante Smith is a highly accomplished university student with an impressive academic record and a strong work ethic. With a CGPA of 8.2, Dante has consistently demonstrated his dedication and commitment to excellence in his studies. He has also completed two internships, highlighting his ability to apply his academic knowledge in a practical setting and gain valuable work experience. Dante has also participated in three projects and two workshops/certifications, demonstrating his versatility and adaptability. In terms of his aptitude, Dante scored 90 on an aptitude test, indicating that he has strong analytical and problem-solving skills. His soft skills rating of 4.5 suggests that he is also well-regarded by his peers and mentors for his professionalism, communication skills, and teamwork abilities. Dante is also actively involved in extracurricular activities, which shows that he values diversity and enjoys engaging with different communities. He has also undergone placement training, indicating that he is eager to apply his skills in a real-world setting and further his career. Overall, Dante Smith is a well-rounded and highly skilled student who has demonstrated his potential to succeed in a variety of settings. With his strong academic record, practical experience, and valuable skills, he is poised to make a significant impact in his field."]
        ],
        title="Placement Buddy",
        description="Here is a placement buddy, which tells your placement changes based on your university records like CGPA, Extra Curricular activities, High school marks, etc.",)
iface.launch()