File size: 2,184 Bytes
4fdacd8 8eb7b79 4fdacd8 178e01a 4fdacd8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
import os
from typing import Optional, List
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from huggingface_hub import HfApi, hf_hub_url
# Configuration
APP_TITLE = "OMol25 API Service"
REPO_ID = "facebook/OMol25"
REPO_TYPE = "model"
# Get token from environment
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
print("WARNING: No HF_TOKEN found. Gated repo access will fail.")
# Initialize API
api = HfApi(token=HF_TOKEN)
app = FastAPI(
title=APP_TITLE,
servers=[{"url": "https://amcgovern-omol25dsf.hf.space"}]
)
# CORS setup for ChatGPT
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=False,
allow_methods=["GET"],
allow_headers=["*"],
)
# Response models
class FileListResponse(BaseModel):
count: int
files: List[str]
class FileURLResponse(BaseModel):
path: str
url: str
# Endpoints
@app.get("/")
def root():
return {"message": f"{APP_TITLE} is running", "endpoints": ["/health", "/files", "/url"]}
@app.get("/health")
def health():
return {"status": "ok", "repo": REPO_ID}
@app.get("/privacy")
def privacy():
return {"privacy": "No user data stored. Public repo access only."}
@app.get("/files", response_model=FileListResponse)
def list_files(q: Optional[str] = Query(None, description="Search filter")):
"""List files in OMol25 repository"""
try:
files = api.list_repo_files(REPO_ID, repo_type=REPO_TYPE)
if q:
files = [f for f in files if q.lower() in f.lower()]
return FileListResponse(count=len(files), files=files[:500])
except Exception as e:
raise HTTPException(status_code=503, detail=f"Cannot access OMol25: {str(e)}")
@app.get("/url", response_model=FileURLResponse)
def get_file_url(path: str = Query(..., description="File path in repo")):
"""Get direct download URL for a file"""
try:
url = hf_hub_url(repo_id=REPO_ID, filename=path, repo_type=REPO_TYPE)
return FileURLResponse(path=path, url=url)
except Exception as e:
raise HTTPException(status_code=400, detail=f"Invalid path: {str(e)}") |