Update app.py

app.py

@@ -66,21 +66,36 @@ def identify_subtopic(question, domain):
 def generate_question(prompt, domain, state):
     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,
-        num_return_sequences=1,
-        no_repeat_ngram_size=2,
-        top_k=30,
+        do_sample=True,
         top_p=0.9,
+        top_k=30,
         temperature=0.7,
-        do_sample=True,
+        no_repeat_ngram_size=2,
         pad_token_id=tokenizer.eos_token_id,
+        eos_token_id=tokenizer.eos_token_id
     )
-    question = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+
+    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()
+
+    # ✅ Ensure it ends with a "?"
     if not question.endswith("?"):
-        question = question.split("?")[0] + "?"
+        question = question.split("?")[0].strip() + "?"
 
+    # ✅ Check for uniqueness
     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)
@@ -88,7 +103,7 @@ def generate_question(prompt, domain, state):
             state["asked_subtopics"].add(subtopic)
         return question
     else:
-        return generate_question(prompt, domain, state)  # Retry
+        return generate_question(prompt, domain, state)  # Try again
 
 def match_company(user_input):
     user_input_lower = user_input.lower()