Create app.py

添加 OPT-125m Gradio 应用

app.py

@@ -0,0 +1,169 @@
+import gradio as gr
+from transformers import pipeline
+import torch
+import time
+
+# --- 配置 ---
+MODEL_ID = "jinv2/opt125m-wikitext2-finetuned"
+TASK = "text-generation"
+
+# --- 设备选择 ---
+# 优先使用 GPU (如果 Space 配置了)
+device = 0 if torch.cuda.is_available() else -1
+device_name = "GPU" if device == 0 else "CPU"
+print(f"使用设备: {device_name}")
+
+# --- 加载模型 Pipeline ---
+# 使用 pipeline 简化文本生成任务
+print(f"开始加载模型: {MODEL_ID}...")
+try:
+    # 对于 OPT 模型，通常不需要 trust_remote_code=True
+    # torch_dtype 设为 'auto' 让 transformers 自动选择最佳精度
+    pipe = pipeline(
+        TASK,
+        model=MODEL_ID,
+        torch_dtype='auto', # 自动选择精度 (float32 on CPU, float16/bfloat16 on GPU if supported)
+        device=device
+    )
+    print("模型加载成功。")
+    # 获取模型实际加载的数据类型
+    if hasattr(pipe.model, 'dtype'):
+         loaded_dtype = pipe.model.dtype
+         print(f"模型加载使用的数据类型: {loaded_dtype}")
+    else:
+         print("无法自动检测模型加载的数据类型，可能使用默认值。")
+
+except Exception as e:
+    print(f"加载模型时出错: {e}")
+    raise gr.Error(f"加载模型 '{MODEL_ID}' 失败。错误: {e}。请检查 Space 日志。")
+
+# --- 文本生成函数 ---
+def generate_text(prompt, max_length, temperature, top_p, repetition_penalty):
+    """使用加载的 pipeline 生成文本"""
+    if not prompt:
+        return "请输入起始文本 (prompt)。"
+
+    print(f"\n收到提示词: '{prompt}'")
+    print(f"生成参数: 最大长度={max_length}, 温度={temperature}, Top-p={top_p}, 重复惩罚={repetition_penalty}")
+
+    # 注意：max_length 通常包含 prompt 的长度。
+    # 我们希望生成 max_new_tokens，所以总长度是 prompt 长度 + max_new_tokens
+    # 但 text-generation pipeline 的 max_length 参数是 *总* 长度。
+    # 为简单起见，我们直接使用 max_length 作为总长度限制，用户输入的 prompt 会被计算在内。
+    # 或者，我们可以计算 prompt 的 token 数量并加上期望的新 token 数。
+    # 这里我们采用更简单的 max_length 方法。
+
+    start_time = time.time()
+    try:
+        # OPT 模型通常用于文本续写，不需要复杂的聊天模板
+        outputs = pipe(
+            prompt,
+            max_length=max_length, # 这是生成的总文本长度，包括 prompt
+            do_sample=True if temperature > 0 else False, # 仅当 temperature > 0 时采样
+            temperature=max(temperature, 1e-6), # Temperature 不能为 0 或负数
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            num_return_sequences=1,
+            pad_token_id=pipe.tokenizer.eos_token_id # 避免填充警告
+        )
+        generated_text = outputs[0]['generated_text']
+
+        # pipeline 输出通常包含原始提示，我们只返回生成的部分
+        # （如果需要完整文本，可以直接返回 generated_text）
+        response = generated_text[len(prompt):].strip()
+
+        end_time = time.time()
+        duration = end_time - start_time
+        print(f"生成完成。原始输出长度: {len(generated_text)}, 提取的续写部分: {response}")
+        print(f"生成耗时: {duration:.2f} 秒")
+
+        # 如果模型有时不生成任何新内容，返回提示信息
+        if not response and len(generated_text) <= len(prompt):
+             return "(模型没有生成新的文本，可能需要调整参数或 prompt)"
+        return response
+
+    except Exception as e:
+        print(f"生成过程中发生错误: {e}")
+        import traceback
+        traceback.print_exc()
+        return f"生成过程中发生错误: {e}"
+
+# --- 创建 Gradio 界面 ---
+with gr.Blocks(theme=gr.themes.Soft(), title=f"测试 {MODEL_ID}") as demo:
+    gr.Markdown(f"""
+    # 测试文本生成模型: `{MODEL_ID}`
+    输入一段起始文本 (prompt)，模型将尝试续写它。
+    **注意:** 模型运行在 **{device_name}** 上。
+    """)
+
+    with gr.Row():
+        with gr.Column(scale=2):
+            prompt_input = gr.Textbox(
+                label="输入起始文本 (Prompt)",
+                lines=5,
+                placeholder="例如：从前有一只勇敢的小兔子，它梦想着..."
+            )
+            with gr.Accordion("高级生成选项", open=False):
+                max_length_slider = gr.Slider(
+                    minimum=20,
+                    maximum=512, # OPT-125m 的标准上下文长度通常是 2048，但设置低一些以防内存问题和过长生成
+                    value=100, # 默认生成较短的续写
+                    step=10,
+                    label="最大总长度 (Max Length)",
+                    info="生成的文本（包括提示）的最大令牌数。"
+                )
+                temperature_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=2.0,
+                    value=0.7,
+                    step=0.05,
+                    label="温度 (Temperature)",
+                    info="控制随机性。>1 更随机, <1 更确定。0 表示贪婪解码。"
+                )
+                top_p_slider = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.9,
+                    step=0.05,
+                    label="Top-p (Nucleus Sampling)",
+                    info="累积概率阈值，用于筛选下一个词的候选。仅在 temperature > 0 时有效。"
+                )
+                repetition_penalty_slider = gr.Slider(
+                    minimum=1.0,
+                    maximum=2.0,
+                    value=1.1,
+                    step=0.1,
+                    label="重复惩罚 (Repetition Penalty)",
+                    info="大于 1 可减少重复。设为 1.0 则禁用。"
+                )
+        submit_button = gr.Button("生成续写", variant="primary")
+
+    with gr.Column(scale=3):
+        output_text = gr.Textbox(
+            label="模型续写内容 (Generated Text)",
+            lines=15,
+            interactive=False
+        )
+
+    gr.Examples(
+        examples=[
+            ["人工智能的未来是", 150, 0.8, 0.9, 1.1],
+            ["今天天气真不错，阳光明媚，", 80, 0.7, 0.95, 1.0],
+            ["The quick brown fox jumps over the", 50, 0.5, 0.9, 1.2],
+        ],
+        inputs=[prompt_input, max_length_slider, temperature_slider, top_p_slider, repetition_penalty_slider],
+        outputs=output_text,
+        fn=generate_text,
+        cache_examples=False,
+        label="示例"
+    )
+
+    submit_button.click(
+        fn=generate_text,
+        inputs=[prompt_input, max_length_slider, temperature_slider, top_p_slider, repetition_penalty_slider],
+        outputs=output_text,
+        api_name="generate"
+    )
+
+# 启动 Gradio 应用
+demo.launch()