Rename app.py to t5_small.py

t5_small.py

@@ -0,0 +1,109 @@
+import streamlit as st
+import torch
+import time
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# 设置页面配置
+st.set_page_config(
+    page_title="文本关键点提取工具",
+    page_icon="📝",
+    layout="wide"
+)
+
+# 标题和介绍
+st.title("文本关键点提取工具")
+st.markdown("基于t5-small模型，从文本中提取关键点")
+
+# 定义模型
+model_list = {
+    "t5-small": "keypoint_T5-Small"
+}
+
+# 缓存模型加载（避免重复加载）
+@st.cache_resource
+def load_model(model_name):
+    st.info(f"正在加载模型: {model_name}")
+    start_time = time.time()
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+    
+    # 判断是否有GPU可用
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device)
+    
+    elapsed = time.time() - start_time
+    st.success(f"✅ 模型加载完成: {model_name}，耗时 {elapsed:.2f} 秒")
+    
+    return model, tokenizer, device
+
+# 生成关键点的函数
+def generate_keypoints(model, tokenizer, device, text, max_new_tokens=64):
+    if not text.strip():
+        return "请输入文本内容"
+    
+    # T5模型的特定提示
+    prompt = f"summarize: {text}"
+    
+    # 编码输入文本
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True).to(device)
+    
+    # 生成关键点
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+    
+    # 解码输出
+    keypoint = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+    
+    # 后处理：规范化"no key point"输出
+    if keypoint.lower() in ["none", "no keypoint", "no key point", "n/a", "na", "", "nothing"]:
+        keypoint = "未提取到关键点"
+        
+    return keypoint
+
+# 主界面
+def main():
+    # 侧边栏
+    with st.sidebar:
+        st.header("模型配置")
+        max_new_tokens = st.slider("最大生成长度", min_value=16, max_value=256, value=64, step=16)
+        
+        # 加载模型
+        model, tokenizer, device = load_model(list(model_list.keys())[0])
+    
+    # 主内容区
+    col1, col2 = st.columns([1, 1])
+    
+    with col1:
+        st.subheader("输入文本")
+        user_text = st.text_area(
+            "请输入需要提取关键点的文本",
+            height=300,
+            placeholder="在此粘贴文本内容..."
+        )
+        
+        if st.button("提取关键点"):
+            if model and tokenizer and device:
+                with st.spinner("正在提取关键点..."):
+                    start_time = time.time()
+                    result = generate_keypoints(model, tokenizer, device, user_text, max_new_tokens)
+                    elapsed = time.time() - start_time
+                    
+                    st.session_state["result"] = result
+                    st.session_state["time"] = elapsed
+                    
+                    st.success(f"✅ 关键点提取完成，耗时 {elapsed:.2f} 秒")
+            else:
+                st.warning("请先确保模型加载成功")
+    
+    with col2:
+        st.subheader("提取结果")
+        if "result" in st.session_state:
+            st.markdown(f"**{list(model_list.values())[0]}:**")
+            st.info(st.session_state["result"])
+            st.caption(f"生成耗时: {st.session_state['time']:.2f} 秒")
+        else:
+            st.info("请输入文本并点击提取按钮")
+
+if __name__ == "__main__":
+    main()