app.py

import json
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import os

# Set up device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

hf_token = os.environ["HF_TOKEN"]


# ===============================
# Load Question Generation Model
# ===============================
model_path = "AI-Mock-Interviewer/T5" 
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

# Move model to the appropriate device
model.to(device)

# ===============================
# Load Evaluation Model (QwQ)
# ===============================
bnb_config = BitsAndBytesConfig(
    load_in_8bit=True,
    llm_int8_enable_fp32_cpu_offload=True,
)

qwq_model_id = "unsloth/QwQ-32B-unsloth-bnb-4bit"
qwq_tokenizer = AutoTokenizer.from_pretrained(qwq_model_id, trust_remote_code=True)
qwq_model = AutoModelForCausalLM.from_pretrained(
    qwq_model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

# ===============================
# Prompts and Scoring
# ===============================
system_prompt = """
You are conducting a mock technical interview. The candidate's experience level can be entry-level, mid-level, or senior-level. Generate questions and follow-up questions based on the domain and the candidate's experience level. Consider these aspects:
1. The question should be relevant to the domain and appropriate for the candidate's experience level.
2. For follow-up questions, analyze the candidate's last response and ask questions that probe deeper into their understanding.
3. Avoid repeating previously asked questions or subtopics.
4. Keep questions clear and concise, targeting core technical and communication skills.
"""

subtopic_keywords = {
    "data analysis": ["data cleaning", "missing data", "EDA", "visualization"],
    "machine learning": ["supervised learning", "overfitting", "hyperparameter tuning"],
    "software engineering": ["code optimization", "design patterns", "database design"],
}

rating_scores = {"Good": 3, "Average": 2, "Needs Improvement": 1}
score_categories = [(90, "Excellent"), (75, "Very Good"), (60, "Good"), (45, "Average"), (0, "Needs Improvement")]

# ===============================
# Helper Functions
# ===============================
def identify_subtopic(question, domain):
    domain = domain.lower()
    if domain in subtopic_keywords:
        for subtopic in subtopic_keywords[domain]:
            if subtopic in question.lower():
                return subtopic
    return None

def generate_question(prompt, domain, state=None):
    full_prompt = system_prompt + "\n" + prompt
    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
    outputs = model.generate(
        inputs["input_ids"],
        max_new_tokens=50,
        no_repeat_ngram_size=2,
        top_k=30,
        top_p=0.9,
        temperature=0.7,
        do_sample=True,
        pad_token_id=tokenizer.eos_token_id,
    )
    question = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
    if not question.endswith("?"):
        question += "?"

    subtopic = identify_subtopic(question, domain)

    if state is not None:
        if (question not in state["asked_questions"] and
            (subtopic is None or subtopic not in state["asked_subtopics"])):
            state["asked_questions"].add(question)
            if subtopic:
                state["asked_subtopics"].add(subtopic)
            return question
    return question

def evaluate_response(response, question):
    eval_prompt = (
        "Evaluate the following candidate response to an interview question.\n\n"
        f"**Question:** {question}\n"
        f"**Candidate's Response:** {response}\n\n"
        "Provide a rating as: 'Good', 'Average', or 'Needs Improvement'.\n"
        "Also, provide a brief suggestion for improvement. Format:\n"
        "Rating: <Rating>\nSuggestion: <Suggestion>"
    )
    inputs = qwq_tokenizer(eval_prompt, return_tensors="pt", padding=True).to(qwq_model.device)
    outputs = qwq_model.generate(
        inputs["input_ids"],
        max_new_tokens=100,
        top_k=30,
        top_p=0.9,
        temperature=0.7,
        do_sample=True,
        pad_token_id=qwq_tokenizer.eos_token_id,
    )
    evaluation = qwq_tokenizer.decode(outputs[0], skip_special_tokens=True)
    rating, suggestion = "Unknown", "No suggestion available."
    for line in evaluation.splitlines():
        if "Rating:" in line:
            rating = line.split("Rating:")[1].strip()
        if "Suggestion:" in line:
            suggestion = line.split("Suggestion:")[1].strip()
    return rating, suggestion

def reset_state(name, domain, company, level):
    return {
        "name": name,
        "domain": domain,
        "company": company,
        "level": level,
        "asked_questions": set(),
        "asked_subtopics": set(),
        "conversation": [],
        "evaluations": [],
        "interview_active": True
    }

def start_interview(name, domain, company, level):
    state = reset_state(name, domain, company, level)
    prompt = f"Domain: {domain}. Candidate experience level: {level}. Generate the first question:"
    question = generate_question(prompt, domain, state)
    state["conversation"].append(("Interviewer", question))
    return state["conversation"], state

def submit_response(response, state):
    if not state["interview_active"]:
        return state["conversation"], state

    if not response.strip():
        state["conversation"].append(("System", "⚠️ Please answer the question before proceeding."))
        return state["conversation"], state

    if response.strip().lower() == "exit":
        return end_interview(state)

    state["conversation"].append(("Candidate", response))
    last_q = [msg for role, msg in reversed(state["conversation"]) if role == "Interviewer"][0]
    rating, suggestion = evaluate_response(response, last_q)

    state["evaluations"].append({
        "question": last_q,
        "response": response,
        "rating": rating,
        "suggestion": suggestion
    })

    state["conversation"].append(("Evaluator", f"Rating: {rating}\nSuggestion: {suggestion}"))
    prompt = f"Domain: {state['domain']}. Candidate's last response: {response}. Generate a follow-up question:"
    follow_up = generate_question(prompt, state["domain"], state)
    state["conversation"].append(("Interviewer", follow_up))
    return state["conversation"], state

def end_interview(state):
    state["interview_active"] = False
    total = sum(rating_scores.get(ev["rating"], 0) for ev in state["evaluations"])
    max_total = len(state["evaluations"]) * 3
    percent = (total / max_total * 100) if max_total > 0 else 0
    category = next(label for threshold, label in score_categories if percent >= threshold)

    summary = {
        "name": state["name"],
        "domain": state["domain"],
        "level": state["level"],
        "company": state["company"],
        "score": f"{total}/{max_total}",
        "percentage": round(percent, 2),
        "category": category,
        "evaluations": state["evaluations"]
    }

    filename = f"sessions/{state['name'].replace(' ', '_').lower()}_session.json"
    os.makedirs("sessions", exist_ok=True)
    with open(filename, "w") as f:
        json.dump(summary, f, indent=4)

    state["conversation"].append(("System", f"✅ Interview ended.\nFinal Score: {summary['score']} ({summary['category']})"))
    return state["conversation"], state

def clear_state():
    return [], None

# ===============================
# Gradio UI
# ===============================
with gr.Blocks() as demo:
    gr.Markdown("# 🧠 AI Mock Interview with Evaluation, History & Exit")

    with gr.Row():
        name_input = gr.Textbox(label="Your Name")
        domain_input = gr.Textbox(label="Domain", placeholder="e.g. Software Engineering")
        company_input = gr.Textbox(label="Company (Optional)", placeholder="e.g. Google")
        level_input = gr.Dropdown(
            label="Experience Level",
            choices=["Entry-Level", "Mid-Level", "Senior-Level"],
            value="Entry-Level"
        )

    start_button = gr.Button("Start Interview")
    chatbot = gr.Chatbot(label="Interview Conversation", height=450)

    with gr.Row():
        response_input = gr.Textbox(label="Your Response (type 'exit' to quit)", lines=2)
        submit_button = gr.Button("Submit")
        exit_button = gr.Button("Exit Interview")
        clear_button = gr.Button("Clear Session")

    state = gr.State()

    start_button.click(start_interview,
                       inputs=[name_input, domain_input, company_input, level_input],
                       outputs=[chatbot, state])
    submit_button.click(submit_response, inputs=[response_input, state], outputs=[chatbot, state]).then(lambda: "", None, response_input)
    exit_button.click(end_interview, inputs=state, outputs=[chatbot, state])
    clear_button.click(lambda: ([], None), outputs=[chatbot, state])

demo.launch()