Update app.py

app.py

@@ -26,8 +26,8 @@ model = AutoModelForCausalLM.from_pretrained(
 model = torch.compile(model)
 
 # System prompt for consistent question generation
-system_prompt = """
-You are conducting a mock technical interview. Your task is to generate clear, concise, and unique interview questions based on the given domain and round. Follow these rules:
+system_prompt = f"""
+You are conducting a {round_type.lower()} interview for a position in {domain} at {company}. Generate one concise and unique question: Follow these rules:
 
 1. Only output one question — do not include explanations, elaborations, or surrounding text.
 2. Do not use any labels like "Follow-up Question" or "Question:" in your output. Just the raw question.
@@ -64,38 +64,34 @@ def identify_subtopic(question, domain):
     return None
 
 def generate_question(prompt, domain, state):
-    full_prompt = system_prompt + "\n" + prompt
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+    full_prompt = f"{system_prompt.strip()}\n{prompt.strip()}"
+    inputs = tokenizer(full_prompt, return_tensors="pt", padding=True, truncation=True).to(device)
 
     outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=50,
-        do_sample=True,
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        max_new_tokens=60,
+        temperature=0.7,
         top_p=0.9,
         top_k=30,
-        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
         no_repeat_ngram_size=2,
-        pad_token_id=tokenizer.eos_token_id,
-        eos_token_id=tokenizer.eos_token_id
     )
 
     decoded = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
-    # ✅ Remove everything before the actual question (i.e., the prompt text)
-    if decoded.startswith(system_prompt.strip()):
-        decoded = decoded[len(system_prompt):].strip()
-    if prompt.strip() in decoded:
-        decoded = decoded.split(prompt.strip())[-1].strip()
-
-    # ✅ Extract only the question line
-    question_lines = decoded.splitlines()
-    question = next((line for line in question_lines if "?" in line), decoded).strip()
+    # Step 1: Strip system prompt and user prompt from output
+    if full_prompt in decoded:
+        decoded = decoded.split(full_prompt)[-1].strip()
 
-    # ✅ Ensure it ends with a "?"
-    if not question.endswith("?"):
-        question = question.split("?")[0].strip() + "?"
+    # Step 2: Take only the first proper sentence ending with '?'
+    lines = decoded.splitlines()
+    question = next((line for line in lines if "?" in line), decoded)
+    question = question.split("?")[0].strip() + "?"
 
-    # ✅ Check for uniqueness
+    # Step 3: Check for duplicates
     subtopic = identify_subtopic(question, domain)
     if question not in state["asked_questions"] and (subtopic is None or subtopic not in state["asked_subtopics"]):
         state["asked_questions"].add(question)
@@ -103,7 +99,7 @@ def generate_question(prompt, domain, state):
             state["asked_subtopics"].add(subtopic)
         return question
     else:
-        return generate_question(prompt, domain, state)  # Try again
+        return generate_question(prompt, domain, state)  # Retry generation
 
 def match_company(user_input):
     user_input_lower = user_input.lower()