Salandra/Experiments

1import os2import requests3from typing import Optional4 5from fastapi import FastAPI, Header, HTTPException, BackgroundTasks6from fastapi.responses import FileResponse7from huggingface_hub.hf_api import HfApi8 9from .models import config, WebhookPayload10 11WEBHOOK_SECRET = os.getenv("WEBHOOK_SECRET")12HF_ACCESS_TOKEN = os.getenv("HF_ACCESS_TOKEN")13AUTOTRAIN_API_URL = "https://api.autotrain.huggingface.co"14AUTOTRAIN_UI_URL = "https://ui.autotrain.huggingface.co"15 16 17app = FastAPI()18 19@app.get("/")20async def home():21	return FileResponse("home.html")22 23@app.post("/webhook")24async def post_webhook(25		payload: WebhookPayload,26		task_queue: BackgroundTasks,27		x_webhook_secret:  Optional[str] = Header(default=None),28	):29	if x_webhook_secret is None:30		raise HTTPException(401)31	if x_webhook_secret != WEBHOOK_SECRET:32		raise HTTPException(403)33	if not (34		payload.event.action == "update"35		and payload.event.scope.startswith("repo.content")36		and payload.repo.name == config.input_dataset37		and payload.repo.type == "dataset"38	):39		# no-op40		return {"processed": False}41 42	task_queue.add_task(43		schedule_retrain,44		payload45	)46 47	return {"processed": True}48 49 50def schedule_retrain(payload: WebhookPayload):51	# Create the autotrain project52	try:53		project = AutoTrain.create_project(payload)54		AutoTrain.add_data(project_id=project["id"])55		AutoTrain.start_processing(project_id=project["id"])56	except requests.HTTPError as err:57		print("ERROR while requesting AutoTrain API:")58		print(f"  code: {err.response.status_code}")59		print(f"  {err.response.json()}")60		raise61	# Notify in the community tab62	notify_success(project["id"])63 64	return {"processed": True}65 66 67class AutoTrain:68	@staticmethod69	def create_project(payload: WebhookPayload) -> dict:70		project_resp = requests.post(71			f"{AUTOTRAIN_API_URL}/projects/create",72			json={73				"username": config.target_namespace,74				"proj_name": f"{config.autotrain_project_prefix}-{payload.repo.headSha[:7]}",75				"task": 18, # image-multi-class-classification76				"config": {77					"hub-model": config.input_model,78					"max_models": 1,79					"language": "unk",80				}81			},82			headers={83				"Authorization": f"Bearer {HF_ACCESS_TOKEN}"84			}85		)86		project_resp.raise_for_status()87		return project_resp.json()88 89	@staticmethod90	def add_data(project_id:int):91		requests.post(92			f"{AUTOTRAIN_API_URL}/projects/{project_id}/data/dataset",93			json={94				"dataset_id": config.input_dataset,95				"dataset_split": "train",96				"split": 4,97				"col_mapping": {98					"image": "image",99					"label": "target",100				}101			},102			headers={103				"Authorization": f"Bearer {HF_ACCESS_TOKEN}",104			}105		).raise_for_status()106 107	@staticmethod108	def start_processing(project_id: int):109		resp = requests.post(110			f"{AUTOTRAIN_API_URL}/projects/{project_id}/data/start_processing",111			headers={112				"Authorization": f"Bearer {HF_ACCESS_TOKEN}",113			}114		)115		resp.raise_for_status()116		return resp117 118 119def notify_success(project_id: int):120	message = NOTIFICATION_TEMPLATE.format(121		input_model=config.input_model,122		input_dataset=config.input_dataset,123		project_id=project_id,124		ui_url=AUTOTRAIN_UI_URL,125	)126	return HfApi(token=HF_ACCESS_TOKEN).create_discussion(127		repo_id=config.input_dataset,128		repo_type="dataset",129		title="✨ Retraining started!",130		description=message,131		token=HF_ACCESS_TOKEN,132	)133 134NOTIFICATION_TEMPLATE = """\135🌸 Hello there!136Following an update of [{input_dataset}](https://huggingface.co/datasets/{input_dataset}), an automatic re-training of [{input_model}](https://huggingface.co/{input_model}) has been scheduled on AutoTrain!137Please review and approve the project [here]({ui_url}/{project_id}/trainings) to start the training job.138(This is an automated message)139"""