app.py

import json
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, BitsAndBytesConfig
import torch
import os
import gradio_client.utils as client_utils
import sys

# ===============================
# Recursion Handling Fix
# ===============================
def _patched_json_schema_to_python_type(schema, defs=None, depth=0):
    # Safety check to prevent infinite recursion
    if depth > 100:
        return "Any"
    
    # Handle boolean cases
    if isinstance(schema, bool):
        return "Any" if schema else "None"
    
    # Call the original function with increased depth
    try:
        return client_utils._json_schema_to_python_type(schema, defs)
    except RecursionError:
        return "Any"

# Modify the utilities to use the patched function
client_utils._json_schema_to_python_type = _patched_json_schema_to_python_type

# Increase recursion limit as a backup
sys.setrecursionlimit(10000)

# 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
    # Explicitly set padding side and add pad token
    tokenizer.padding_side = "left"
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Tokenize with explicit padding and attention mask
    inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to(device)
    
    outputs = model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        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.pad_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"].append(question)
            if subtopic:
                state["asked_subtopics"].append(subtopic)
            return question
    return question


def evaluate_response(response, question):
    # Explicitly set padding side and add pad token
    qwq_tokenizer.padding_side = "left"
    if qwq_tokenizer.pad_token is None:
        qwq_tokenizer.pad_token = qwq_tokenizer.eos_token
    
    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>"
    )
    
    # Tokenize with explicit padding and attention mask
    inputs = qwq_tokenizer(eval_prompt, return_tensors="pt", padding=True, truncation=True).to(qwq_model.device)
    
    outputs = qwq_model.generate(
        inputs["input_ids"],
        attention_mask=inputs["attention_mask"],
        max_new_tokens=100,
        top_k=30,
        top_p=0.9,
        temperature=0.7,
        do_sample=True,
        pad_token_id=qwq_tokenizer.pad_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": [],
        "asked_subtopics": [],
        "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({"role": "Interviewer", "content": 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({"role": "System", "content": "⚠️ Please answer the question before proceeding."})
        return state["conversation"], state

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

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

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

    state["conversation"].append({"role": "Evaluator", "content": 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({"role": "Interviewer", "content": 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({"role": "System", "content": f"✅ Interview ended. \nFinal Score: {summary['score']} ({summary['category']})"})
    return state["conversation"], state

def clear_state():
    return [], reset_state("", "", "", "Entry-Level")

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

    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, type="messages")

    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")

    # Initialize state with proper structure
    state = gr.State(value=reset_state("", "", "", "Entry-Level"))

    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(clear_state, outputs=[chatbot, state])

demo.launch()