update

app.py

@@ -1,14 +1,95 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+import os
+from dotenv import load_dotenv
+
+# 导入 utils 模块
+from utils.chat_request import ChatRequest
+from utils.chat_response import create_chat_response, ChatResponse
+from utils.model import check_model, initialize_pipeline, download_model, DownloadRequest
+
+# 全局变量
+model_name = None
+pipe = None
+tokenizer = None
 
 # 初始化 FastAPI 应用
 app = FastAPI(title="HF-Model-Runner API", version="0.0.1")
 
-model_name = None
+@app.on_event("startup")
+async def startup_event():
+    """
+    应用启动时初始化 pipeline
+    """
+    global pipe, tokenizer, model_name
+    
+    # 加载 .env 文件
+    load_dotenv()
+    
+    # 从 .env 获取默认模型名称，如果没有则使用默认值
+    default_model = os.getenv("DEFAULT_MODEL_NAME", "unsloth/functiongemma-270m-it")
+    print(f"应用启动，正在初始化模型: {default_model}")
+    
+    try:
+        pipe, tokenizer, success = initialize_pipeline(default_model)
+        if success:
+            model_name = default_model
+            print(f"✓ 模型 {default_model} 初始化成功")
+        else:
+            print(f"✗ 模型 {default_model} 初始化失败")
+    except Exception as e:
+        print(f"✗ 启动时模型初始化失败: {e}")
 
 @app.get("/")
-def greet_json():
-    return {
-        "message": "HF-Model-Runner API is running!",
-        "model": model_name,
-        "status": "ready"
-    }
+async def read_root():
+    return {"message": "Welcome to HF-Model-Runner API! Visit /docs for API documentation."}
+
+@app.post("/download")
+async def download_model_endpoint(request: DownloadRequest):
+    """
+    下载指定的 HuggingFace 模型
+    """
+    global pipe, tokenizer, model_name
+    
+    try:
+        success, message = download_model(request.model)
+        if success:
+            # 下载成功后，直接初始化该模型
+            pipe, tokenizer, init_success = initialize_pipeline(request.model)
+            if init_success:
+                model_name = request.model
+                return {
+                    "status": "success", 
+                    "message": message,
+                    "loaded": True,
+                    "current_model": model_name
+                }
+            else:
+                return {
+                    "status": "success", 
+                    "message": message,
+                    "loaded": False,
+                    "error": "模型下载成功但初始化失败"
+                }
+        else:
+            raise HTTPException(status_code=500, detail=message)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/v1/chat/completions", response_model=ChatResponse)
+async def chat_completions(request: ChatRequest):
+    """
+    OpenAI 兼容的聊天完成接口
+    """
+    global pipe, tokenizer, model_name
+    
+    # 检查模型是否匹配，如果请求的模型与当前加载的模型不同，需要重新初始化
+    if request.model != model_name:
+        pipe, tokenizer, success = initialize_pipeline(request.model)
+        if not success:
+            raise HTTPException(status_code=500, detail="模型初始化失败")
+        model_name = request.model
+    
+    try:
+        return create_chat_response(request, pipe, tokenizer)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))

requirements.txt

@@ -5,3 +5,4 @@ huggingface_hub
 torch
 accelerate
 python-multipart
+python-dotenv

utils/chat_request.py

@@ -0,0 +1,12 @@
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
+
+
+class ChatRequest(BaseModel):
+    model: Optional[str] = "gpt-3.5-turbo"
+    messages: List[Dict[str, Any]]
+    temperature: Optional[float] = 1.0
+    max_tokens: Optional[int] = None
+    top_p: Optional[float] = 1.0
+    frequency_penalty: Optional[float] = 0.0
+    presence_penalty: Optional[float] = 0.0

utils/chat_response.py

@@ -0,0 +1,115 @@
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
+import time
+import re
+
+# 聊天响应模型
+class ChatChoice(BaseModel):
+    index: int
+    message: Dict[str, str]
+    finish_reason: str
+
+
+class ChatUsage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class ChatResponse(BaseModel):
+    id: str
+    object: str
+    created: int
+    model: str
+    choices: List[ChatChoice]
+    usage: ChatUsage
+
+
+def convert_json_format(input_data):
+    """转换 pipeline 输出格式"""
+    output_generations = []
+    for item in input_data:
+        generated_text_list = item.get('generated_text', [])
+        
+        assistant_content = ""
+        for message in generated_text_list:
+            if message.get('role') == 'assistant':
+                assistant_content = message.get('content', '')
+                break
+
+        # 移除 </think>...</think> 标签
+        clean_content = re.sub(r'\s*', '', assistant_content, flags=re.DOTALL).strip()
+
+        output_generations.append([
+            {
+                "text": clean_content,
+                "generationInfo": {
+                    "finish_reason": "stop"
+                }
+            }
+        ])
+    
+    return {"generations": output_generations}
+
+
+def create_chat_response(request: Any, pipe=None, tokenizer=None) -> ChatResponse:
+    """
+    创建聊天响应 - 使用 pipeline 生成实际响应
+    """
+    if pipe is None:
+        # 如果 pipeline 未初始化，返回模拟响应
+        response_message = {
+            "role": "assistant",
+            "content": "模型正在初始化中，请稍后重试..."
+        }
+        completion_text = response_message["content"]
+    else:
+        # 使用 pipeline 生成响应
+        messages = request.messages
+        
+        # 从 request 获取 max_new_tokens，如果没有则使用默认值 1000
+        # max_new_tokens = request.max_tokens if request.max_tokens is not None else 1000
+        max_new_tokens = request.max_tokens if request.max_tokens is not None else None
+        
+        # 调用 pipeline
+        result = pipe(messages, max_new_tokens=max_new_tokens)
+        # result = pipe(messages)
+        
+        # 转换格式
+        converted_result = convert_json_format(result)
+        
+        # 获取生成的文本
+        completion_text = converted_result["generations"][0][0]["text"]
+        
+        response_message = {
+            "role": "assistant",
+            "content": completion_text
+        }
+    
+    # 计算 token 数量
+    if tokenizer:
+        prompt_tokens = sum(len(tokenizer.encode(msg.get("content", ""))) for msg in request.messages)
+        completion_tokens = len(tokenizer.encode(completion_text))
+    else:
+        # 简化估算
+        prompt_tokens = sum(len(msg.get("content", "")) for msg in request.messages) // 4
+        completion_tokens = len(completion_text) // 4
+    
+    return ChatResponse(
+        id=f"chatcmpl-{int(time.time())}",
+        object="chat.completion",
+        created=int(time.time()),
+        model=request.model,
+        choices=[
+            ChatChoice(
+                index=0,
+                message=response_message,
+                finish_reason="stop"
+            )
+        ],
+        usage=ChatUsage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens
+        )
+    )

utils/model.py

@@ -0,0 +1,105 @@
+import os
+import sys
+from pathlib import Path
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import login
+from fastapi import HTTPException
+from pydantic import BaseModel
+
+
+class DownloadRequest(BaseModel):
+    model: str
+
+
+def check_model(model_name):
+    """
+    检查模型是否存在
+    参数: model_name - 从 request 传递过来的模型名称
+    返回: (model_name, cache_dir, success)
+    """
+    cache_dir = "./my_model_cache"
+    
+    # 检查模型是否已存在于缓存中
+    model_path = Path(cache_dir) / f"models--{model_name.replace('/', '--')}"
+    snapshot_path = model_path / "snapshots"
+    
+    if snapshot_path.exists() and any(snapshot_path.iterdir()):
+        print(f"✓ 模型 {model_name} 已存在于缓存中")
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+            return model_name, cache_dir, True
+        except Exception as e:
+            print(f"⚠ 加载现有模型失败: {e}")
+            return model_name, cache_dir, False
+    else:
+        raise HTTPException(status_code=404, detail=f"模型 `{model_name}` 不存在，请先下载")
+
+
+def download_model(model_name):
+    """
+    下载指定的模型
+    参数: model_name - 要下载的模型名称
+    返回: (success, message)
+    """
+    cache_dir = "./my_model_cache"
+    
+    print(f"开始下载模型: {model_name}")
+    print(f"缓存目录: {cache_dir}")
+    
+    # 登录 Hugging Face（可选，用于需要认证的模型）
+    token = os.getenv("HUGGINGFACE_TOKEN")
+    if token:
+        try:
+            print("登录 Hugging Face...")
+            login(token=token)
+            print("✓ HuggingFace 登录成功！")
+        except Exception as e:
+            print(f"⚠ 登录失败: {e}")
+            print("继续使用公开模型")
+    else:
+        print("ℹ 未设置 HUGGINGFACE_TOKEN - 仅使用公开模型")
+    
+    try:
+        # 下载 tokenizer
+        print("正在下载 tokenizer...")
+        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+        print("✓ Tokenizer 下载成功！")
+        
+        # 下载模型
+        print("正在下载模型...")
+        model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir)
+        print("✓ 模型下载成功！")
+        
+        print(f"✓ 模型和 tokenizer 已成功下载到 {cache_dir}")
+        return True, f"模型 {model_name} 下载成功"
+        
+    except Exception as e:
+        print(f"✗ 下载模型时出错: {e}")
+        return False, f"下载失败: {str(e)}"
+
+
+def initialize_pipeline(model_name):
+    """
+    使用模型初始化 pipeline
+    参数: model_name - 从 request 传递过来的模型名称
+    返回: (pipe, tokenizer, success)
+    """
+    model_name, cache_dir, success = check_model(model_name)
+    
+    if not success:
+        return None, None, False
+    
+    try:
+        # 确保 tokenizer 已加载
+        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
+        
+        print(f"使用 {model_name} 初始化 pipeline...")
+        # 移除 cache_dir 参数，只传递给 AutoTokenizer 和 AutoModelForCausalLM
+        pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)
+        print("✓ Pipeline 初始化成功！")
+        
+        return pipe, tokenizer, True
+        
+    except Exception as e:
+        print(f"✗ Pipeline 初始化失败: {e}")
+        return None, None, False