utils/bio_generator.py

"""
Adoption bio generator using Hugging Face models.
Primary: HF Inference API (cloud). Fallback: local flan-t5-base.
"""

import os
import logging

logger = logging.getLogger(__name__)

PROMPT_TEMPLATE = """You are a creative writer for an animal shelter. Write a warm, engaging adoption bio (1-2 paragraphs, about 100-150 words) for the following animal.

Name: {name}
Type: {animal_type}
Traits: {traits}

The bio should highlight the animal's personality, make potential adopters feel an emotional connection, and end with a gentle call to action encouraging adoption. Keep the tone upbeat and heartfelt. Do not use hashtags. Write only the bio text, nothing else."""


def _build_prompt(name: str, animal_type: str, characteristics: list[str]) -> str:
    if not name:
        name = "This sweet friend"
    traits = ", ".join(characteristics) if characteristics else "friendly and lovable"
    return PROMPT_TEMPLATE.format(name=name, animal_type=animal_type, traits=traits)


def generate_bio_api(name: str, animal_type: str, characteristics: list[str]) -> str | None:
    token = os.environ.get("HF_TOKEN")
    if not token:
        logger.info("No HF_TOKEN set, skipping API generation")
        return None

    try:
        from huggingface_hub import InferenceClient
        client = InferenceClient(token=token)
        prompt = _build_prompt(name, animal_type, characteristics)

        response = client.text_generation(
            prompt,
            model="mistralai/Mistral-7B-Instruct-v0.3",
            max_new_tokens=300,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.2,
        )
        return response.strip()
    except Exception as e:
        logger.warning(f"API bio generation failed: {e}")
        return None


def generate_bio_local(name: str, animal_type: str, characteristics: list[str]) -> str:
    try:
        from transformers import pipeline

        generator = pipeline(
            "text2text-generation",
            model="google/flan-t5-base",
            max_new_tokens=250,
        )
        prompt = _build_prompt(name, animal_type, characteristics)
        result = generator(prompt, max_new_tokens=250)
        text = result[0]["generated_text"].strip()
        if text and len(text) > 80:
            return text
    except Exception as e:
        logger.warning(f"Local model bio generation failed: {e}")

    return _generate_template_bio(name, animal_type, characteristics)


def _generate_template_bio(name: str, animal_type: str, characteristics: list[str]) -> str:
    if not name:
        name = "This adorable friend"

    trait_text = ""
    if characteristics:
        if len(characteristics) == 1:
            trait_text = characteristics[0].lower()
        elif len(characteristics) == 2:
            trait_text = f"{characteristics[0].lower()} and {characteristics[1].lower()}"
        else:
            trait_text = ", ".join(c.lower() for c in characteristics[:-1]) + f", and {characteristics[-1].lower()}"

    type_desc = {
        "Dog": "pup",
        "Cat": "kitty",
        "Other": "little friend",
    }.get(animal_type, "companion")

    paragraphs = []

    opening_options = [
        f"Meet {name}, a wonderful {type_desc} who is ready to find a forever home!",
        f"Say hello to {name}! This lovable {type_desc} has been waiting patiently for the perfect family.",
        f"Allow us to introduce {name}, an absolutely charming {type_desc} with a heart full of love to give.",
    ]
    import random
    paragraphs.append(random.choice(opening_options))

    if trait_text:
        body = f"{name} is {trait_text}. "
    else:
        body = f"{name} has a wonderful personality. "

    type_activities = {
        "Dog": "Whether it's a walk in the park, a game of fetch, or just curling up on the couch together, "
               f"{name} will be your most loyal companion.",
        "Cat": f"Whether it's a quiet afternoon nap in a sunbeam or an evening play session, "
               f"{name} knows how to make every moment special.",
        "Other": f"With a unique personality that will brighten your days, "
                 f"{name} is sure to bring joy and wonder into your life.",
    }
    body += type_activities.get(animal_type, type_activities["Other"])
    paragraphs.append(body)

    closing = (
        f"Could {name} be the missing piece in your family? "
        f"Visit the shelter today and let {name} show you just how much love one {type_desc} can bring. "
        f"Every animal deserves a home, and {name} is ready to make yours complete."
    )
    paragraphs.append(closing)

    return "\n\n".join(paragraphs)


def generate_bio(name: str, animal_type: str, characteristics: list[str]) -> str:
    bio = generate_bio_api(name, animal_type, characteristics)
    if bio:
        return bio

    return generate_bio_local(name, animal_type, characteristics)