Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 import spaces
 import torch
+import difflib
 from threading import Thread
 from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM, TextIteratorStreamer
 
@@ -34,6 +35,30 @@ def load_model():
             
     return model, tokenizer
 
+def create_diff_html(original, cleaned):
+    """Create HTML diff visualization"""
+    original_lines = original.splitlines(keepends=True)
+    cleaned_lines = cleaned.splitlines(keepends=True)
+    
+    differ = difflib.unified_diff(original_lines, cleaned_lines, fromfile='Original', tofile='Cleaned', lineterm='')
+    
+    html_diff = '<div style="font-family: monospace; font-size: 12px; white-space: pre-wrap;">'
+    
+    for line in differ:
+        if line.startswith('+++') or line.startswith('---'):
+            html_diff += f'<div style="color: #666;">{line}</div>'
+        elif line.startswith('@@'):
+            html_diff += f'<div style="color: #0066cc; font-weight: bold;">{line}</div>'
+        elif line.startswith('+'):
+            html_diff += f'<div style="background-color: #e6ffed; color: #24292e;">{line}</div>'
+        elif line.startswith('-'):
+            html_diff += f'<div style="background-color: #ffeef0; color: #24292e;">{line}</div>'
+        else:
+            html_diff += f'<div>{line}</div>'
+    
+    html_diff += '</div>'
+    return html_diff
+
 @spaces.GPU(duration=60)
 def clean_text(text):
     model, tokenizer = load_model()
@@ -65,28 +90,53 @@ def clean_text(text):
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
     
-    # Yield text as it's generated
+    # Collect generated text, skipping the input
     generated_text = ""
+    input_length = len(formatted_text)
+    full_output = ""
+    
     for new_text in streamer:
-        generated_text += new_text
-        yield generated_text
+        full_output += new_text
+        # Only yield the part after the input
+        if len(full_output) > input_length:
+            generated_text = full_output[input_length:].strip()
+            yield generated_text, ""
     
     thread.join()
+    
+    # After generation is complete, create diff
+    diff_html = create_diff_html(text, generated_text)
+    yield generated_text, diff_html
 
-iface = gr.Interface(
-    fn=clean_text,
-    inputs=gr.Textbox(
-        lines=5,
-        placeholder="Enter text to clean...",
-        label="Input Text"
-    ),
-    outputs=gr.Textbox(
-        lines=5,
-        label="Cleaned Text"
-    ),
-    title="TextClean-4B Demo",
-    description="Simple demo for text cleaning using textcleanlm/textclean-4B model"
-)
+# Create the interface with blocks for better control
+with gr.Blocks(title="TextClean-4B Demo") as demo:
+    gr.Markdown("# TextClean-4B Demo")
+    gr.Markdown("Simple demo for text cleaning using textcleanlm/textclean-4B model")
+    
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                lines=5,
+                placeholder="Enter text to clean...",
+                label="Input Text"
+            )
+            submit_btn = gr.Button("Clean Text", variant="primary")
+    
+    with gr.Row():
+        output_text = gr.Textbox(
+            lines=5,
+            label="Cleaned Text",
+            interactive=False
+        )
+    
+    with gr.Row():
+        diff_display = gr.HTML(label="Diff View")
+    
+    submit_btn.click(
+        fn=clean_text,
+        inputs=input_text,
+        outputs=[output_text, diff_display]
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()