Update app.py

app.py

@@ -7,17 +7,50 @@ from huggingface_hub import InferenceClient  # add to requirements.txt
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 YOUTUBE_RE = re.compile(r"https?://(?:www\.)?youtube\.com/watch\?v=[\w-]+")
 
+NUM_WORDS = {
+    "zero":"0","one":"1","two":"2","three":"3","four":"4","five":"5",
+    "six":"6","seven":"7","eight":"8","nine":"9","ten":"10","eleven":"11",
+    "twelve":"12","thirteen":"13","fourteen":"14","fifteen":"15","sixteen":"16",
+    "seventeen":"17","eighteen":"18","nineteen":"19","twenty":"20"
+}
+
+def _extract_bare_number(text: str) -> str | None:
+    """Return the first number found as a string (prefers integers, falls back to decimals or number-words)."""
+    line = text.strip().splitlines()[0]
+
+    # 1) integer
+    m = re.search(r"(?<![\d.])[-+]?\d+(?![\d.])", line)
+    if m:
+        return m.group(0).lstrip("+")
+
+    # 2) decimal (if ever needed)
+    m = re.search(r"[-+]?\d+\.\d+", line)
+    if m:
+        return m.group(0).lstrip("+")
+
+    # 3) number words → digits
+    mw = re.search(r"\b(" + "|".join(NUM_WORDS.keys()) + r")\b", line.lower())
+    if mw:
+        return NUM_WORDS[mw.group(1)]
+
+    return None
+
 def format_final_answer(q: str, raw: str) -> str:
-    text = raw.strip().splitlines()[0]
-    if "how many" in q.lower():
-        m = re.search(r"\d+", text)
-        if m:
-            return m.group(0)
-    # As a safeguard, also strip common wrappers:
+    text = raw.strip()
     for pre in ("final answer:", "answer:", "final:", "prediction:"):
         if text.lower().startswith(pre):
             text = text[len(pre):].strip()
-    return text
+            break
+    
+    # If the question implies a numeric answer, force a bare number
+    ql = question.lower()
+    if any(k in ql for k in ["how many", "number", "highest number", "count", "total", "included"]):
+        n = _extract_bare_number(text)
+        if n is not None:
+            return n  # <-- always a string, e.g. "3"
+    
+    # otherwise, keep first line as-is (already stripped)
+    return text.splitlines()[0]
 
 # --- provider selection (HF serverless text-generation by default; optional Groq) ---
 def select_model():