Create app.py

app.py

@@ -0,0 +1,70 @@
+import gradio as gr
+import spaces
+import torch
+from transformers import AutoModel, AutoTokenizer, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+
+model_id = "textcleanlm/textclean-4B"
+model = None
+tokenizer = None
+
+def load_model():
+    global model, tokenizer
+    if model is None:
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        
+        # Add padding token if needed
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Try different model classes
+        for model_class in [AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoModel]:
+            try:
+                model = model_class.from_pretrained(
+                    model_id,
+                    torch_dtype=torch.bfloat16,
+                    device_map="auto"
+                )
+                break
+            except:
+                continue
+                
+        if model is None:
+            raise ValueError(f"Could not load model {model_id}")
+            
+    return model, tokenizer
+
+@spaces.GPU(duration=60)
+def clean_text(text):
+    model, tokenizer = load_model()
+    
+    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True)
+    inputs = {k: v.cuda() for k, v in inputs.items()}
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=512,
+            num_beams=4,
+            early_stopping=True
+        )
+    
+    cleaned_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return cleaned_text
+
+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"
+)
+
+if __name__ == "__main__":
+    iface.launch()