test the model as well

app.py

@@ -195,10 +195,19 @@ def extract_content(html, model, tokenizer):
     return f"{extracted_content} (confidence: {confidence:.2f})"
 
 
-def gradio_interface(html_input):
+def test_model(html_input, model, tokenizer):
+    result = extract_content(html_input, model, tokenizer)
+    return result
+
+
+def gradio_interface(html_input, test_type):
     global trained_model, trained_tokenizer
-    extracted_content = extract_content(html_input, trained_model, trained_tokenizer)
-    return extracted_content
+    if test_type == "Custom Input":
+        return test_model(html_input, trained_model, trained_tokenizer)
+    elif test_type == "Generate Random HTML":
+        random_html, _ = generate_html_content()
+        result = test_model(random_html, trained_model, trained_tokenizer)
+        return random_html, result
 
 
 print("Starting training process...")
@@ -207,10 +216,16 @@ print("Training completed. Launching Gradio interface...")
 
 iface = gr.Interface(
     fn=gradio_interface,
-    inputs=gr.Textbox(lines=10, label="Input HTML"),
-    outputs=gr.Textbox(label="Extracted Content"),
+    inputs=[
+        gr.Textbox(lines=10, label="Input HTML"),
+        gr.Radio(["Custom Input", "Generate Random HTML"], label="Test Type", value="Custom Input")
+    ],
+    outputs=[
+        gr.Textbox(lines=10, label="HTML Content"),
+        gr.Textbox(label="Extracted Information")
+    ],
     title="HTML Content Extractor",
-    description="Enter HTML content to extract information."
+    description="Enter HTML content or generate random HTML to test the model."
 )
 
 iface.launch()