Update app.py

app.py

@@ -19,8 +19,8 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 print("Loading local model...")
 
 generator = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base",
+    "text-generation",
+    model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
     device_map="auto"
 )
 
@@ -61,9 +61,11 @@ Answer:
 
             text = result[0]["generated_text"]
 
-            answer = text.strip()
+            if "Answer:" in text:
+                answer = text.split("Answer:")[-1].strip()
+            else:
+                answer = text.strip()
 
-            # cleanup
             answer = answer.split("\n")[0].strip()
 
             return answer
@@ -207,7 +209,7 @@ Answer:
 
         if "what year" in q:
 
-            years = re.findall(r'\b(19|20)\d{2}\b', question)
+            years = re.findall(r'\b(?:19|20)\d{2}\b', question)
 
             if years:
                 return years[0]