Spaces:

kota101
/

ta_ml

Sleeping

ta_ml / api /tokenize.py

muhammadjalaludin

Implement service layers for backup, model management, file handling, tokenization, and training

1424118 over 1 year ago

4.53 kB

# ta_ml/api/tokenize.py
"""
API endpoints for text tokenization.
"""
import logging
from typing import Optional

import pandas as pd
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
from fastapi.responses import StreamingResponse

from api.dependencies import get_tokenization_service
from data.file_handler import FileHandler, FileHandlerError
from models.api import (TokenizationDatasetResponse,
                        TokenizationResult)
from services.tokenization_service import TokenizationService

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/tokenize", tags=["Tokenization"])


@router.get(
    "",
    response_model=TokenizationResult,
    summary="Tokenize a single word or phrase",
    description="Tokenizes a single string using the model's tokenizer and returns detailed information, including the sub-word tokens, their corresponding IDs, and their vocabulary status.",
    responses={
        200: {"description": "Successful tokenization."},
        422: {"description": "The 'word' query parameter is missing or empty."},
    },
)
def tokenize_word(
    word: str = Query(..., min_length=1,
                      description="The word or phrase to tokenize."),
    tokenization_service: TokenizationService = Depends(
        get_tokenization_service),
):
    return tokenization_service.tokenize_word(word)


@router.post(
    "/file",
    summary="Tokenize a text column from a file",
    description=(
        "Tokenizes a text column from an uploaded CSV or Excel file. "
        "This endpoint has two modes:\n\n"
        "1.  **JSON Response (default):** Returns detailed statistics and tokenization results for a slice of the data (controlled by `offset` and `limit`).\n"
        "2.  **File Export:** If `export_format` is set to `csv` or `xlsx`, the endpoint processes the data slice and returns the full tokenization results as a downloadable file."
    ),
    responses={
        200: {
            "description": "Successfully tokenized the data. The response will be either a JSON object or a file stream, depending on the `export_format` parameter.",
            "content": {
                "application/json": {
                    "schema": {
                        "$ref": "#/components/schemas/TokenizationDatasetResponse"
                    }
                },
                "text/csv": {},
                "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": {},
            },
        },
        400: {
            "description": "Bad request due to issues like unsupported file format, text column not found, or no data available for the given offset/limit."
        },
        422: {"description": "Validation error, e.g., the file was not included in the form data."},
    },
)
async def tokenize_file(
    file: UploadFile = File(..., description="A CSV or Excel file."),
    text_column: str = Query(
        "text", description="The name of the column containing text to tokenize."),
    limit: int = Query(100, ge=1, le=10000,
                       description="Maximum number of rows to process."),
    offset: int = Query(
        0, ge=0, description="Number of rows to skip from the beginning."),
    export_format: Optional[str] = Query(None, enum=[
                                         "csv", "xlsx"], description="If provided, returns the tokenized data as a downloadable file."),
    tokenization_service: TokenizationService = Depends(
        get_tokenization_service),
):
    try:
        df = FileHandler.read_to_dataframe(file)
    except FileHandlerError as e:
        raise HTTPException(status_code=400, detail=str(e))

    if text_column not in df.columns:
        raise HTTPException(
            status_code=400, detail=f"Column '{text_column}' not found.")

    df_slice = df.iloc[offset: offset + limit]
    if df_slice.empty:
        raise HTTPException(
            status_code=400, detail="No data found for the given offset and limit.")

    results, stats = tokenization_service.tokenize_dataframe(
        df_slice, text_column)

    if export_format:
        output_df = pd.DataFrame([res.model_dump() for res in results])
        buffer, media_type, filename = FileHandler.export_dataframe(
            output_df, export_format)
        headers = {
            "Content-Disposition": f"attachment; filename=tokenized_{filename}"}
        return StreamingResponse(buffer, media_type=media_type, headers=headers)

    return TokenizationDatasetResponse(stats=stats, results=results)

	# ta_ml/api/tokenize.py
	"""
	API endpoints for text tokenization.
	"""
	import logging
	from typing import Optional

	import pandas as pd
	from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
	from fastapi.responses import StreamingResponse

	from api.dependencies import get_tokenization_service
	from data.file_handler import FileHandler, FileHandlerError
	from models.api import (TokenizationDatasetResponse,
	TokenizationResult)
	from services.tokenization_service import TokenizationService

	logger = logging.getLogger(__name__)
	router = APIRouter(prefix="/tokenize", tags=["Tokenization"])


	@router.get(
	"",
	response_model=TokenizationResult,
	summary="Tokenize a single word or phrase",
	description="Tokenizes a single string using the model's tokenizer and returns detailed information, including the sub-word tokens, their corresponding IDs, and their vocabulary status.",
	responses={
	200: {"description": "Successful tokenization."},
	422: {"description": "The 'word' query parameter is missing or empty."},
	},
	)
	def tokenize_word(
	word: str = Query(..., min_length=1,
	description="The word or phrase to tokenize."),
	tokenization_service: TokenizationService = Depends(
	get_tokenization_service),
	):
	return tokenization_service.tokenize_word(word)


	@router.post(
	"/file",
	summary="Tokenize a text column from a file",
	description=(
	"Tokenizes a text column from an uploaded CSV or Excel file. "
	"This endpoint has two modes:\n\n"
	"1. JSON Response (default): Returns detailed statistics and tokenization results for a slice of the data (controlled by `offset` and `limit`).\n"
	"2. File Export: If `export_format` is set to `csv` or `xlsx`, the endpoint processes the data slice and returns the full tokenization results as a downloadable file."
	),
	responses={
	200: {
	"description": "Successfully tokenized the data. The response will be either a JSON object or a file stream, depending on the `export_format` parameter.",
	"content": {
	"application/json": {
	"schema": {
	"$ref": "#/components/schemas/TokenizationDatasetResponse"
	}
	},
	"text/csv": {},
	"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": {},
	},
	},
	400: {
	"description": "Bad request due to issues like unsupported file format, text column not found, or no data available for the given offset/limit."
	},
	422: {"description": "Validation error, e.g., the file was not included in the form data."},
	},
	)
	async def tokenize_file(
	file: UploadFile = File(..., description="A CSV or Excel file."),
	text_column: str = Query(
	"text", description="The name of the column containing text to tokenize."),
	limit: int = Query(100, ge=1, le=10000,
	description="Maximum number of rows to process."),
	offset: int = Query(
	0, ge=0, description="Number of rows to skip from the beginning."),
	export_format: Optional[str] = Query(None, enum=[
	"csv", "xlsx"], description="If provided, returns the tokenized data as a downloadable file."),
	tokenization_service: TokenizationService = Depends(
	get_tokenization_service),
	):
	try:
	df = FileHandler.read_to_dataframe(file)
	except FileHandlerError as e:
	raise HTTPException(status_code=400, detail=str(e))

	if text_column not in df.columns:
	raise HTTPException(
	status_code=400, detail=f"Column '{text_column}' not found.")

	df_slice = df.iloc[offset: offset + limit]
	if df_slice.empty:
	raise HTTPException(
	status_code=400, detail="No data found for the given offset and limit.")

	results, stats = tokenization_service.tokenize_dataframe(
	df_slice, text_column)

	if export_format:
	output_df = pd.DataFrame([res.model_dump() for res in results])
	buffer, media_type, filename = FileHandler.export_dataframe(
	output_df, export_format)
	headers = {
	"Content-Disposition": f"attachment; filename=tokenized_{filename}"}
	return StreamingResponse(buffer, media_type=media_type, headers=headers)

	return TokenizationDatasetResponse(stats=stats, results=results)