Update app.py

app.py

@@ -1,5 +1,6 @@
 import torch
 from torch import Tensor, nn
+import gradio as gr
 from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
 
 
@@ -42,9 +43,24 @@ def load_t5(device: str | torch.device = "cuda", max_length: int = 512) -> HFEmb
     # max length 64, 128, 256 and 512 should work (if your sequence is short enough)
     return HFEmbedder("city96/t5-v1_1-xxl-encoder-bf16", max_length=max_length, torch_dtype=torch.bfloat16).to(device)
 
+def run_t5_and_save(text):
+    try:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"使用设备: {device}")
+        t5 = load_t5(device, max_length=512)
+        embeddings = t5([text])
+        output_path = "/home/user/embeddings.pt"
+        torch.save(embeddings, output_path)
+        return f"嵌入形状: {embeddings.shape}", output_path
+    except Exception as e:
+        return f"运行错误: {e}", None
+
 if __name__ == "__main__":
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    t5 = load_t5(device, max_length=256)
-    texts = ["这是一个测试句子。"]
-    embeddings = t5(texts)
-    print(f"Embeddings shape: {embeddings.shape}")
+    iface = gr.Interface(
+        fn=run_t5_and_save,
+        inputs=gr.Textbox(label="输入文本"),
+        outputs=[gr.Textbox(label="结果"), gr.File(label="下载嵌入文件")],
+        title="T5 Embedder",
+        description="输入文本，生成 T5 嵌入并保存为文件"
+    )
+    iface.launch()