Update app.py

app.py

@@ -25,7 +25,18 @@ def extract_info(sentence):
         headers={"Authorization": f"Bearer {HF_API_TOKEN}"},
         json={"inputs": prompt}
     )
-    return response.json()
+    result = response.json()
+    generated_text = result[0]["generated_text"]
+
+    # Try to extract JSON object from the generated text
+    match = re.search(r'\{.*\}', generated_text, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group())
+        except json.JSONDecodeError:
+            return {"error": "Failed to parse JSON"}
+    else:
+        return {"error": "No JSON found in model response"}
 
 # Gradio UI
 iface = gr.Interface(