app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from transformers import pipeline
import torch

# Initialize the FastAPI app
app = FastAPI()

# Determine device (use GPU if available, otherwise CPU)
device = 0 if torch.cuda.is_available() else -1

# Initialize the NER pipeline
ner_pipeline = pipeline(
    "ner",
    model="dbmdz/bert-large-cased-finetuned-conll03-english",
    aggregation_strategy="simple",  # Updated to replace deprecated grouped_entities
    device=device
)

# Initialize the QA pipeline
qa_pipeline = pipeline(
    "question-answering",
    model="deepset/roberta-base-squad2",
    device=device
)

# Allowed domains for filtering
allowed_domains = [
    "clothing", "fashion", "shopping", "accessories", "sustainability", "shoes", "hats", "shirts",
    "dresses", "pants", "jeans", "skirts", "jackets", "coats", "t-shirts", "sweaters", "hoodies",
    "activewear", "formal wear", "casual wear", "sportswear", "outerwear", "swimwear", "underwear",
    "lingerie", "socks", "scarves", "gloves", "belts", "ties", "caps", "beanies", "boots", "sandals",
    "heels", "sneakers", "materials", "cotton", "polyester", "wool", "silk", "leather", "denim",
    "linen", "athleisure", "ethnic wear", "fashion trends", "custom clothing", "tailoring",
    "sustainable materials", "recycled clothing", "fashion brands", "streetwear",
    "footwear", "handbags", "jewelry", "watches", "eyewear", "cosmetics", "beauty products",
    "personal care", "fragrances", "home decor", "lifestyle", "luxury goods", "vintage clothing",
    "second-hand clothing", "upcycled fashion", "ethical fashion", "eco-friendly products",
    "fashion technology", "textile innovation", "fashion marketing", "fashion retail"
]

# Pydantic models for structured response
class Entity(BaseModel):
    word: str
    entity_group: str
    score: float

class NERResponse(BaseModel):
    entities: List[Entity]

class QAResponse(BaseModel):
    question: str
    answer: str
    score: float

class CombinedRequest(BaseModel):
    text: str  # The input text prompt

class CombinedResponse(BaseModel):
    ner: NERResponse  # NER output
    qa: QAResponse  # QA output

# Function to check if the input text belongs to allowed domains
def is_text_in_allowed_domain(text: str, domains: List[str]) -> bool:
    for domain in domains:
        if domain in text.lower():
            return True
    return False

# Combined endpoint for NER and QA with domain filtering
@app.post("/process/", response_model=CombinedResponse)
async def process_request(request: CombinedRequest):
    """
    Process the input text for both NER and QA, returning both responses,
    only if the text matches the allowed domains.
    """
    context_msg = "Hutter Products GmbH provides a wide array of services to help businesses create high-quality, sustainable products. Their offerings include comprehensive product design, ensuring items are both visually appealing and functional, and product consulting, which provides expert advice on features, materials, and design elements. They also offer sustainability consulting to integrate eco-friendly practices, such as using recycled materials and Ocean Bound Plastic. Additionally, they manage customized production to ensure products meet the highest standards and offer product animation services, creating realistic rendered images and animations to enhance online engagement. These services collectively enable businesses to develop products that are sustainable, market-responsive, and aligned with their brand identity."
    input_text = request.text

    # Check if the input text belongs to the allowed domains
    if not is_text_in_allowed_domain(input_text, allowed_domains):
        raise HTTPException(
            status_code=400,
            detail=(
                "The input text does not match the allowed domains. "
                "Please provide a query related to clothing, fashion, or accessories."
            )
        )

    # Perform Named Entity Recognition (NER)
    ner_entities = ner_pipeline(input_text)

    # Process the NER entities into the required format
    formatted_entities = [
        {
            "word": entity["word"],
            "entity_group": entity["entity_group"],
            "score": float(entity["score"]),  # Convert numpy.float32 to Python float
        }
        for entity in ner_entities
    ]
    ner_response = {"entities": formatted_entities}

    # Perform Question Answering (QA)
    qa_result = qa_pipeline(question=input_text, context=context_msg)
    qa_result["score"] = float(qa_result["score"])  # Convert numpy.float32 to Python float

    qa_response = {
        "question": input_text,
        "answer": qa_result["answer"],
        "score": qa_result["score"]
    }

    # Return both NER and QA responses
    return {"ner": ner_response, "qa": qa_response}

# Root endpoint
@app.get("/")
async def root():
    """
    Root endpoint to confirm the server is running.
    """
    return {"message": "Welcome to the filtered NER and QA API!"}