Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from peft import PeftModel | |
| # --- 模型加载配置 --- | |
| ADAPTER_REPO_ID = "jinv2/gpt2-lora-trajectory-prediction" | |
| BASE_MODEL_NAME = "gpt2" | |
| # --- 加载模型和分词器 --- | |
| # 这是一个耗时操作，Gradio应用启动时会执行一次 | |
| print(f"开始加载模型: {BASE_MODEL_NAME} 和适配器: {ADAPTER_REPO_ID}") | |
| try: | |
| base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_NAME) | |
| tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO_ID) # 适配器仓库通常包含分词器配置 | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| print("tokenizer.pad_token 设置为 tokenizer.eos_token") | |
| model = PeftModel.from_pretrained(base_model, ADAPTER_REPO_ID) | |
| model.eval() # 设置为评估模式 | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model.to(device) | |
| print(f"模型和分词器加载完成，运行在: {device}") | |
| except Exception as e: | |
| print(f"模型加载失败: {e}") | |
| # 如果模型加载失败，Gradio界面可能无法正常工作，这里可以抛出错误或设置一个标志 | |
| model = None | |
| tokenizer = None | |
| raise RuntimeError(f"无法加载模型: {e}") | |
| # --- 推理函数 --- | |
| def predict_trajectory(history_text_input): | |
| if model is None or tokenizer is None: | |
| return "错误: 模型未能成功加载，请检查Space的日志。" | |
| if not history_text_input or not history_text_input.strip(): | |
| return "错误: 请输入有效的历史轨迹。" | |
| # 格式化为模型期望的输入 | |
| # 假设用户只输入历史点，例如 "1.00,1.00,0.50,0.00; 1.05,1.00,0.50,0.00" | |
| prompt = f"历史: {history_text_input.strip()}; 预测:" | |
| print(f"收到的提示: {prompt}") | |
| try: | |
| inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device) | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=60, # 调整以适应预期的输出长度 (例如2-3个点) | |
| num_return_sequences=1, | |
| pad_token_id=tokenizer.pad_token_id, # 使用pad_token_id | |
| eos_token_id=tokenizer.eos_token_id, | |
| # temperature=0.7, # 如果想要一些随机性 | |
| # do_sample=True, # 如果想要一些随机性 | |
| do_sample=False, # 为了演示的确定性 | |
| num_beams=1 # 使用贪婪解码 | |
| ) | |
| generated_text_full = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| print(f"完整生成文本: {generated_text_full}") | |
| predicted_part = "" | |
| # 从完整输出中提取预测部分 | |
| if "预测:" in generated_text_full: | |
| split_output = generated_text_full.split("预测:", 1) | |
| if len(split_output) > 1: | |
| predicted_part = split_output[1].strip() | |
| # 清理可能的末尾分号或eos token的文本残留 | |
| if predicted_part.endswith(tokenizer.eos_token): | |
| predicted_part = predicted_part[:-len(tokenizer.eos_token)].strip() | |
| if predicted_part.endswith(';'): | |
| predicted_part = predicted_part[:-1].strip() | |
| else: | |
| # 如果模型输出不包含 "预测:"，尝试从提示后截取 | |
| # 这部分逻辑可能需要根据模型的实际输出行为调整 | |
| if prompt in generated_text_full: | |
| predicted_part = generated_text_full[len(prompt):].strip() | |
| else: # 假设模型只输出了预测部分 (可能需要更鲁棒的逻辑) | |
| predicted_part = generated_text_full.strip() # 基本回退 | |
| if predicted_part.endswith(tokenizer.eos_token): | |
| predicted_part = predicted_part[:-len(tokenizer.eos_token)].strip() | |
| if predicted_part.endswith(';'): | |
| predicted_part = predicted_part[:-1].strip() | |
| print(f"提取的预测部分: {predicted_part}") | |
| return predicted_part | |
| except Exception as e: | |
| print(f"推理时发生错误: {e}") | |
| import traceback | |
| traceback.print_exc() | |
| return f"推理错误: {str(e)}" | |
| # --- 创建 Gradio 界面 --- | |
| # 使用 gr.Markdown 来显示更丰富的文本和说明 | |
| readme_url = f"https://huggingface.co/{ADAPTER_REPO_ID}" | |
| description = f""" | |
| # GPT-2 LoRA 轨迹预测 Demo | |
| 这是一个使用微调后的 `gpt2` 模型进行轨迹预测的简单演示。 | |
| 模型仓库: [{ADAPTER_REPO_ID}]({readme_url}) | |
| **如何使用:** | |
| 1. 在下面的 "历史轨迹输入" 框中输入历史轨迹点。 | |
| 2. 格式应为: `x1,y1,vx1,vy1; x2,y2,vx2,vy2` (例如，两个历史点，用分号隔开)。 | |
| 3. 每个点包含四个逗号分隔的数值: x坐标, y坐标, x方向速度, y方向速度。 | |
| 4. 点击 "预测轨迹" 按钮查看模型生成的未来轨迹点。 | |
| """ | |
| # 示例输入 | |
| example_history = "0.00,0.00,1.00,0.00; 0.10,0.00,1.00,0.00" | |
| # 定义界面组件 | |
| iface = gr.Interface( | |
| fn=predict_trajectory, | |
| inputs=gr.Textbox( | |
| lines=3, | |
| placeholder="例如: 0.00,0.00,1.00,0.00; 0.10,0.00,1.00,0.00", | |
| label="历史轨迹输入 (格式: x1,y1,vx1,vy1; x2,y2,vx2,vy2; ...)", | |
| value=example_history # 设置一个默认示例值 | |
| ), | |
| outputs=gr.Textbox( | |
| lines=3, | |
| label="模型预测的未来轨迹 (文本格式)" | |
| ), | |
| title="基于LLM的轨迹预测", | |
| description=description, | |
| examples=[ | |
| ["1.00,1.00,0.50,0.00; 1.05,1.00,0.50,0.00"], | |
| ["-2.0,0.5,0.0,1.0; -2.0,0.6,0.0,1.0; -2.0,0.7,0.0,1.0"], # 三个历史点 | |
| ["0.0,0.0,0.2,0.2; 0.02,0.02,0.2,0.2"] | |
| ], | |
| allow_flagging='never' # 通常用于演示，不需要用户标记 | |
| ) | |
| # 启动 Gradio 应用 (在 Hugging Face Spaces 上会自动处理) | |
| if __name__ == "__main__": | |
| if model is not None and tokenizer is not None: # 仅当模型加载成功时启动 | |
| print("正在本地启动Gradio应用...") | |
| iface.launch() | |
| else: | |
| print("模型未能加载，Gradio应用无法启动。请检查日志。") |