Added system prompt

app.py

@@ -1,73 +1,142 @@
 import gradio as gr
 from transformers import T5Tokenizer, T5ForConditionalGeneration
+import torch
+from transformers import AutoTokenizer
+
+system_prompt = """
+You are conducting a mock technical interview. Generate questions and follow-up questions based on the domain provided. Consider these aspects:
+1. The question should be relevant to the domain (e.g., software engineering, machine learning).
+2. For follow-up questions, analyze the candidate's last response and ask questions that probe deeper into their understanding, challenge their approach, or request clarification.
+3. The follow-up question should aim to explore the candidate's depth of knowledge and ability to adapt.
+4. If no clear follow-up can be derived, generate a fresh, related question.
+Important: Ensure that each question is clear, concise, and allows the candidate to demonstrate their technical and communicative abilities effectively.
+"""
 
-# Load your Hugging Face model
 model_name = "vai0511/flan-t5-ai-mock-interviewer"
 
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
 
-def init_interview(domain, company):
-    prompt = f"Generate first interview question for {domain} at {company}:"
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    first_question = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    return [
-        (None, f"Interview started for {domain} at {company or 'a company'}"),
-        (None, first_question)
-    ]
+# Initialize conversation state
+def init_state():
+    return {
+        "domain": "",
+        "company": "",
+        "questions": [],
+        "current_step": 0,
+        "max_questions": 10
+    }
 
-def respond(history, domain, company):
-    last_user_answer = history[-1][0]  # Get user's last input
+def generate_question(prompt):
+    full_prompt = system_prompt + "\n" + prompt
+    inputs = tokenizer(full_prompt, return_tensors="pt", return_attention_mask=False).to("cuda")
     
-    prompt = f"""
-    You are a technical interviewer for {company or 'a tech company'} 
-    assessing a candidate in {domain}. 
-    
-    The candidate just answered: {last_user_answer}
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=100,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.7,
+        pad_token_id=tokenizer.eos_token_id,
+    )
     
-    Generate a follow-up question that:
-    1. Probes deeper into technical knowledge
-    2. Relates to their previous answer
-    3. Focuses on {domain} concepts
+    question = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    question = question.replace(full_prompt, "").strip()
     
-    Next question:
-    """
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    next_question = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if not question.endswith("?"):
+        question = question.split("?")[0] + "?"
     
-    return history + [(None, next_question)]
+    return question
 
-with gr.Blocks() as demo:
-    gr.Markdown("## AI Interview Practice")
+def respond(history, message, state):
+    if state["current_step"] == 0:
+        # First question generation
+        prompt = f"Domain: {state['domain']}. "
+        if state["company"]:
+            prompt += f"Company: {state['company']}. "
+        prompt += "Generate the first technical question:"
+        
+        question = generate_question(prompt)
+        state["questions"].append({"question": question, "response": None})
+        state["current_step"] += 1
+        
+        return [
+            (None, f"Starting {state['domain']} interview for {state['company'] or 'general'}"),
+            (None, question)
+        ], state
     
-    with gr.Row():
-        domain = gr.Textbox(label="Domain")
-        company = gr.Textbox(label="Company (Optional)")
+    else:
+        # Store response and generate follow-up
+        state["questions"][-1]["response"] = message
+        last_response = message
+        
+        prompt = f"Domain: {state['domain']}. Candidate's last response: {last_response}. Generate follow-up:"
+        question = generate_question(prompt)
+        
+        # If no good follow-up, generate new question
+        if "?" not in question:
+            prompt = f"Domain: {state['domain']}. Generate new question:"
+            question = generate_question(prompt)
+        
+        state["questions"].append({"question": question, "response": None})
+        state["current_step"] += 1
+        
+        history.append((message, None))  # User response
+        history.append((None, question))  # AI question
+        return history, state
+
+with gr.Blocks(title="AI Interviewer") as demo:
+    state = gr.State(init_state)
     
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Your Answer")
-    start_btn = gr.Button("Start Interview")
-    submit_btn = gr.Button("Submit Answer")
+    with gr.Column():
+        gr.Markdown("## Technical Interview Simulator")
+        
+        with gr.Row():
+            domain = gr.Textbox(label="Domain (e.g., Machine Learning)", interactive=True)
+            company = gr.Textbox(label="Company (Optional)", interactive=True)
+            start_btn = gr.Button("Start Interview")
+        
+        chatbot = gr.Chatbot(label="Conversation")
+        msg = gr.Textbox(label="Your Answer", interactive=True)
+        submit_btn = gr.Button("Submit Answer")
+        end_btn = gr.Button("End Interview")
 
+    # Start interview
     start_btn.click(
-        init_interview,
-        inputs=[domain, company],
-        outputs=[chatbot]
+        fn=lambda d, c, s: (s.update({"domain": d, "company": c, "current_step": 0}), s),
+        inputs=[domain, company, state],
+        outputs=[state]
+    ).then(
+        respond,
+        inputs=[chatbot, gr.Textbox("", visible=False), state],
+        outputs=[chatbot, state]
     )
-    
+
+    # Handle responses
     submit_btn.click(
-        lambda message: [(message, None)],
+        lambda x: x,  # Temporary to capture input
         inputs=[msg],
-        outputs=[chatbot],
-        queue=False
+        outputs=[msg]
     ).then(
         respond,
-        inputs=[chatbot, domain, company],
-        outputs=[chatbot]
-    ).then(lambda: "", None, msg)
+        inputs=[chatbot, msg, state],
+        outputs=[chatbot, state]
+    ).then(
+        lambda: "",  # Clear input
+        None,
+        msg
+    )
+
+    # End interview
+    end_btn.click(
+        lambda: [(None, "Interview ended. Thank you!")],
+        None,
+        chatbot
+    ).then(
+        init_state,
+        None,
+        state
+    )
 
 demo.launch()