File size: 9,038 Bytes
fe9092a
8035330
97809c3
253bbca
fe9092a
6ac4ade
8035330
 
fe9092a
3b22eeb
fe9092a
3b22eeb
2d83410
fe9092a
becbbb8
fe9092a
 
716758b
2d83410
fe9092a
c892caa
c0cfc84
fe9092a
716758b
 
 
253bbca
fe9092a
253bbca
716758b
 
 
 
 
 
 
 
 
 
 
cc9c5eb
fe9092a
 
419fdd6
fe9092a
419fdd6
8035330
bc828c5
fe9092a
 
e51c8cf
1a058f6
 
 
5ee18ab
0a15754
ec5b5a2
1a058f6
 
 
6ac4ade
 
 
 
51484a8
5ee18ab
 
 
 
 
 
 
 
 
e51c8cf
 
 
6ac4ade
 
 
148a6b2
fe9092a
e51c8cf
 
 
 
419fdd6
 
e51c8cf
 
 
 
 
 
 
 
8035330
 
 
ef69b25
 
 
 
fe9092a
ef69b25
8035330
 
c0cfc84
 
 
 
 
3b22eeb
 
 
becbbb8
3b22eeb
 
 
fe9092a
2d83410
 
 
 
 
75141de
 
8035330
3b22eeb
becbbb8
75141de
becbbb8
3b22eeb
75141de
97809c3
2d83410
8035330
 
 
 
 
253bbca
8035330
 
419fdd6
fe9092a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
646eba0
fe9092a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4c88070
810d3a0
2be9bc2
 
4c88070
810d3a0
 
4c88070
 
 
810d3a0
fe9092a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
import asyncio
import json
import logging
import os
import re
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Annotated, List

from cashews import NOT_NONE, cache
from dotenv import load_dotenv
from fastapi import BackgroundTasks, FastAPI, Header, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from httpx import AsyncClient
from huggingface_hub import CommitScheduler, DatasetCard, HfApi, hf_hub_download, whoami
from huggingface_hub.utils import disable_progress_bars
from huggingface_hub.utils._errors import HTTPError
from langfuse.openai import AsyncOpenAI  # OpenAI integration
from pydantic import BaseModel, Field
from starlette.responses import RedirectResponse
from card_processing import parse_markdown, try_load_text, is_empty_template
import logging
import logging.config
import yaml

disable_progress_bars()
load_dotenv()

# Load the logging configuration from the YAML file
with open("log_conf.yaml", "r") as f:
    log_config = yaml.safe_load(f.read())
    logging.config.dictConfig(log_config)

# Create a logger for your application
logger = logging.getLogger("app")

# Use the logger in your application
logger.info("Application started")

Gb = 1073741824
cache.setup("disk://", size_limit=16 * Gb)  # configure as in-memory cache

VOTES_FILE = "data/votes.jsonl"

HF_TOKEN = os.getenv("HF_TOKEN")

hf_api = HfApi(token=HF_TOKEN)
async_httpx_client = AsyncClient()

scheduler = CommitScheduler(
    repo_id="davanstrien/summary-ratings",
    repo_type="dataset",
    folder_path="data",
    path_in_repo="data",
    every=5,
    token=HF_TOKEN,
    hf_api=hf_api,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Running startup event")
    if not Path(VOTES_FILE).exists():
        path = hf_hub_download(
            repo_id="davanstrien/summary-ratings",
            filename="data/votes.jsonl",
            repo_type="dataset",
            token=HF_TOKEN,
            local_dir=".",
            local_dir_use_symlinks=False,
        )
        logger.info(f"Downloaded votes.jsonl to {path}")
    else:
        logger.info("Votes file already exists")
    yield


app = FastAPI(lifespan=lifespan)
# Configure CORS
# origins = [
#     "https://huggingface.co/*",
#     # "chrome-extension://deckahggoiaphiebdipfbiinmaihfpbk",  # Replace with your Chrome plugin ID
# ]


# # Configure CORS settings
# app.add_middleware(
#     CORSMiddleware,
#     allow_origins=["https://huggingface.co/*"],  # Update with your frontend URL
#     allow_credentials=True,
#     allow_methods=["*"],
#     allow_headers=["*"],
# )


def save_vote(vote_entry):
    with scheduler.lock:
        with open(VOTES_FILE, "a") as file:
            date_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            vote_entry["timestamp"] = date_time
            file.write(json.dumps(vote_entry) + "\n")
        logger.info(f"Vote saved: {vote_entry}")


@app.get("/", include_in_schema=False)
def root():
    return RedirectResponse(url="/docs")


class Vote(BaseModel):
    dataset: str
    description: str
    vote: int = Field(..., ge=-1, le=1)
    userID: str


def validate_token(token: str = Header(None)) -> bool:
    try:
        whoami(token)
        return True
    except HTTPError:
        return False


@app.post("/vote")
async def receive_vote(
    vote: Vote,
    Authorization: Annotated[str, Header()],
    background_tasks: BackgroundTasks,
):
    if not validate_token(Authorization):
        logger.error("Invalid token")
        raise HTTPException(status_code=401, detail="Invalid token")
    vote_entry = {
        "dataset": vote.dataset,
        "vote": vote.vote,
        "description": vote.description,
        "userID": vote.userID,
    }
    # Append the vote entry to the JSONL file
    background_tasks.add_task(save_vote, vote_entry)
    return JSONResponse(content={"message": "Vote submitted successfully"})


def format_prompt(card: str) -> str:
    return f"""
Write a tl;dr summary of a dataset based on the dataset card. Focus on the most critical aspects of the dataset. 
The summary should aim to concisely describe the dataset.

CARD: \n\n{card[:6000]}
---

\n\nInstructions:
If the card provides the necessary information, say what the dataset can be used for.
You do not need to mention that the dataset is hosted or available on the Hugging Face Hub.
Do not mention the license of the dataset.
Do not mention the number of examples in the training or test split.
Only mention size if there is extensive discussion of the scale of the dataset in the dataset card.
Do not speculate on anything not explicitly mentioned in the dataset card.
In general avoid references to the quality of the dataset i.e. don't use phrases like 'a high-quality dataset' in the summary.

\n\nOne sentence summary:"""


async def check_when_dataset_last_modified(dataset_id: str) -> datetime | None:
    try:
        response = await async_httpx_client.get(
            f"https://huggingface.co/api/datasets/{dataset_id}"
        )
        if last_modified := response.json().get("lastModified"):
            return datetime.fromisoformat(last_modified)
        return None
    except Exception as e:
        logger.error(e)
        return None


@cache(ttl="48h", condition=NOT_NONE, key="predict:{dataset_id}")
async def predict(card: str, dataset_id: str) -> str | None:
    try:
        prompt = format_prompt(card)
        client = AsyncOpenAI(
            base_url="https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1/v1",
            api_key=HF_TOKEN,
        )

        chat_completion = await client.chat.completions.create(
            model="tgi",
            messages=[
                {"role": "user", "content": prompt},
            ],
            stream=False,
            tags=["tldr-summaries"],
        )
        return chat_completion.choices[0].message.content.strip()
    except Exception as e:
        logger.error(e)
        return None


@app.get("/summary")
async def get_summary(dataset_id: str) -> str | None:
    """
    Get a summary for a dataset based on the provided dataset ID.

    Args:
        dataset_id (str): The ID of the dataset to retrieve the summary for.

    Returns:
        str | None: The generated summary for the dataset, or None if no summary is available or an error occurs."""

    try:
        # dataset_id = request.dataset_id
        try:
            resp = await async_httpx_client.get(
                f"https://huggingface.co/datasets/{dataset_id}/raw/main/README.md",
                follow_redirects=True,
            )
            if resp.status_code != 200:
                return None
        except Exception as e:
            logger.error(e)
            return None
        card_text = resp.text
        card = DatasetCard(card_text)
        text = card.text
        parsed_text = parse_markdown(text)
        if is_empty_template(parsed_text):
            return None
        cache_key = f"predict:{dataset_id}"
        cached_data = await cache.get(cache_key)
        if cached_data is not None:
            cached_summary, cached_last_modified_time = cached_data
            # Get the current last modified time of the dataset
            current_last_modified_time = await check_when_dataset_last_modified(
                dataset_id
            )

            if (
                current_last_modified_time is None
                or cached_last_modified_time >= current_last_modified_time
            ):
                # Use the cached summary if the cached last modified time is greater than or equal to the current last modified time
                logger.info("Using cached summary")
                return cached_summary
        summary = await predict(parsed_text, dataset_id)
        current_last_modified_time = await check_when_dataset_last_modified(dataset_id)
        await cache.set(cache_key, (summary, current_last_modified_time))
        return summary
    except Exception as e:
        logger.error(e)
        return None


class SummariesRequest(BaseModel):
    dataset_ids: List[str]


@cache(ttl="1h", condition=NOT_NONE)
@app.post("/summaries")
async def get_summaries(request: SummariesRequest) -> dict:
    """
    Get summaries for a list of datasets based on the provided dataset IDs.

    Args:
        dataset_ids (List[str]): A list of dataset IDs to retrieve the summaries for.

    Returns:
        dict: A dictionary mapping dataset IDs to their corresponding summaries.
    """
    dataset_ids = request.dataset_ids

    async def get_summary_wrapper(dataset_id):
        return dataset_id, await get_summary(dataset_id)

    summary_tasks = [get_summary_wrapper(dataset_id) for dataset_id in dataset_ids]
    summaries = dict(await asyncio.gather(*summary_tasks))
    for dataset_id in dataset_ids:
        if summaries[dataset_id] is None:
            summaries[dataset_id] = "No summary available"
    return summaries