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):
    try:
        # Print all input parameters
        print(f"Start Interview Called:")
        print(f"Name: {name}")
        print(f"Domain: {domain}")
        print(f"Company: {company}")
        print(f"Level: {level}")

        # Validate inputs
        if not name or not domain:
            return [{"role": "System", "content": "Please provide a name and domain"}], None

        # Create initial state
        state = reset_state(name, domain, company, level)
        print("State reset successfully")

        # Prepare prompt for question generation
        prompt = f"Domain: {domain}. Candidate experience level: {level}. Generate the first question:"
        print(f"Prompt for question generation: {prompt}")

        # Verify model is ready
        print("Model device:", model.device)
        print("Model ready:", model is not None)

        # Generate first question
        try:
            question = generate_question(prompt, domain, state)
            print(f"Generated Question: {question}")
        except Exception as q_error:
            print(f"Question Generation Error: {q_error}")
            question = f"Error generating question: {q_error}"

        # Append question to conversation
        state["conversation"].append({"role": "Interviewer", "content": question})
        
        return state["conversation"], state

    except Exception as e:
        print(f"CRITICAL ERROR in start_interview: {e}")
        import traceback
        traceback.print_exc()
        return [{"role": "System", "content": f"Critical error: {e}"}], None

    print("Conversation returned to UI:", state["conversation"])

def submit_response(response, state):
    print("Submit Response Called")
    print("Response:", response)
    print("Current State:", state)

    # Ensure state is not None and interview is active
    if state is None:
        print("State is None, resetting")
        state = reset_state("", "", "", "Entry-Level")
    
    if not state.get("interview_active", False):
        print("Interview not active")
        return state["conversation"], state

    # Handle empty response
    if not response or not response.strip():
        print("Empty response")
        state["conversation"].append({"role": "System", "content": "⚠️ Please answer the question before proceeding."})
        return state["conversation"], state

    # Exit condition
    if response.strip().lower() == "exit":
        print("Exit requested")
        return end_interview(state)

    # Add candidate response to conversation
    state["conversation"].append({"role": "Candidate", "content": response})
    
    # Find the last interviewer question
    last_q = next((msg["content"] for msg in reversed(state["conversation"]) if msg["role"] == "Interviewer"), "")
    
    # Evaluate response
    print("Evaluating response to question:", last_q)
    rating, suggestion = evaluate_response(response, last_q)

    # Add evaluation to conversation and state
    state["evaluations"].append({
        "question": last_q,
        "response": response,
        "rating": rating,
        "suggestion": suggestion
    })

    state["conversation"].append({"role": "Evaluator", "content": f"Rating: {rating}\nSuggestion: {suggestion}"})

    # Generate follow-up question
    prompt = f"Domain: {state['domain']}. Candidate's last response: {response}. Generate a follow-up question:"
    follow_up = generate_question(prompt, state["domain"], state)
    
    print("Generated Follow-up Question:", follow_up)
    state["conversation"].append({"role": "Interviewer", "content": follow_up})

    return state["conversation"], state

    print("Conversation returned to UI:", state["conversation"])

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