Update app.py

app.py

@@ -2,65 +2,55 @@ import gradio as gr
 from transformers import pipeline
 import spacy
 from textblob import TextBlob
-import requests
 from gradio_client import Client
 
 # Initialize models
-nlp = spacy.load("en_core_web_sm")  # More accurate NER
+nlp = spacy.load("en_core_web_sm")  # NER model
 spell_checker = pipeline("text2text-generation", model="oliverguhr/spelling-correction-english-base")
 
-def preprocess_and_forward(text: str) -> str:
-    """
-    Processes the input text and forwards it to the Gradio client for space_17.
-    Returns only the final translated text.
-    """
-    processed_text, _ = preprocess_text(text)
-    return forward_to_translation(processed_text)
-
 def preprocess_text(text: str):
     """
     Applies spell-checking and named entity recognition (NER) to preprocess text.
+    Returns token-level suggestions.
     """
-    result = {
-        "spell_suggestions": [],
-        "entities": [],
-        "tags": []
-    }
-
-    # Basic spell checking using TextBlob
-    corrected_text = str(TextBlob(text).correct())
-    if corrected_text != text:
-        result["spell_suggestions"].append({
-            "original": text,
-            "corrected": corrected_text
-        })
-
-    # Transformer-based spell check
-    spell_checked = spell_checker(text, max_length=512)[0]['generated_text']
-    if spell_checked != text and spell_checked != corrected_text:
-        result["spell_suggestions"].append({
-            "original": text,
-            "corrected": spell_checked
-        })
-
-    # NER with spaCy
+    tokens = text.split()
+    suggestions = []
+    
+    for token in tokens:
+        token_suggestions = {"original": token, "suggestions": []}
+        
+        # Basic spell checking
+        corrected = str(TextBlob(token).correct())
+        if corrected != token:
+            token_suggestions["suggestions"].append(corrected)
+        
+        # Transformer-based spell checking
+        spell_checked = spell_checker(token, max_length=20)[0]['generated_text']
+        if spell_checked != token and spell_checked not in token_suggestions["suggestions"]:
+            token_suggestions["suggestions"].append(spell_checked)
+        
+        suggestions.append(token_suggestions)
+    
+    # Named Entity Recognition (NER)
     doc = nlp(text)
-    result["entities"] = [{"text": ent.text, "label": ent.label_} for ent in doc.ents]
+    entities = [{"text": ent.text, "label": ent.label_} for ent in doc.ents]
+    
+    return {"tokens": suggestions, "entities": entities}
 
-    # Extract potential tags (hashtags, mentions, etc.)
-    result["tags"] = [token.text for token in doc if token.text.startswith(('#', '@'))]
-
-    # Choose the best-corrected version
-    final_text = spell_checked if spell_checked != text else corrected_text
-
-    return final_text, result
+def preprocess_and_forward(text: str):
+    """
+    Processes the input text, returns suggestions, and forwards the cleaned version for translation.
+    """
+    processed_data = preprocess_text(text)
+    final_text = " ".join([t['suggestions'][0] if t['suggestions'] else t['original'] for t in processed_data["tokens"]])
+    translation = forward_to_translation(final_text)
+    return {"suggestions": processed_data, "translation": translation}
 
 def forward_to_translation(text: str) -> str:
     """
-    Sends preprocessed text to space_17 for translation and returns only the translated text.
+    Sends preprocessed text for translation and returns only the translated text.
     """
     client = Client("Frenchizer/space_17")
-
     try:
         return client.predict(text)
     except Exception as e:
@@ -69,9 +59,11 @@ def forward_to_translation(text: str) -> str:
 # Gradio interface
 with gr.Blocks() as demo:
     input_text = gr.Textbox(label="Input Text")
-    output_text = gr.Textbox(label="Output Text")  # Returns only text
+    output_text = gr.Textbox(label="Translated Text")
+    suggestion_output = gr.JSON(label="Suggestions")
+    
     preprocess_button = gr.Button("Process")
-    preprocess_button.click(fn=preprocess_and_forward, inputs=[input_text], outputs=[output_text])
+    preprocess_button.click(fn=preprocess_and_forward, inputs=[input_text], outputs=[suggestion_output, output_text])
 
 if __name__ == "__main__":
     demo.launch()