Files
2024-02-13 18:02:16 -05:00

91 lines
2.9 KiB
Python

from typing import List
from openai import OpenAI
from pydantic import BaseModel, Field
client = OpenAI()
class SearchQuery(BaseModel):
product_name: str
query: str = Field(
...,
description="A descriptive query to search for the product, include adjectives, and the product type. will be used to serve relevant products to the user.",
)
class MultiSearchQuery(BaseModel):
products: List[SearchQuery]
def extract_table(url: str):
completion = client.chat.completions.create(
model="gpt-4-vision-preview",
max_tokens=1800,
temperature=0,
stop=["```"],
messages=[
{
"role": "system",
"content": f"""
You are an expert system designed to extract products from images for a ecommerse application
Please provide the product name and a descriptive query to search for the product.
Accuratly identify every product in an image and provide a descriptive query to search for the product
You just return a correctly formatted JSON object with the product name and query for each product in the image
and follows the schema below:
{MultiSearchQuery.model_json_schema()}
""",
},
{
"role": "user",
"content": [
{
"type": "text",
"text": "Extract the products from the image, and describe them in a query in JSON format",
},
{
"type": "image_url",
"image_url": {"url": url},
},
],
},
{
"role": "assistant",
"content": "Here is the following search queries for the products in the image\n ```json",
},
],
)
return MultiSearchQuery.model_validate_json(completion.choices[0].message.content)
if __name__ == "__main__":
url = "https://mensfashionpostingcom.files.wordpress.com/2020/03/fbe79-img_5052.jpg?w=768"
products = extract_table(url)
print(products.model_dump_json(indent=2))
"""
{
"products": [
{
"product_name": "Olive Green Shirt",
"query": "Olive green casual long sleeve button-down shirt"
},
{
"product_name": "Black Jeans",
"query": "Slim fit black jeans for men"
},
{
"product_name": "Sunglasses",
"query": "Classic brown aviator sunglasses"
},
{
"product_name": "Leather Strap Watch",
"query": "Minimalist men's watch with black leather strap"
},
{
"product_name": "Beige Sneakers",
"query": "Men's beige lace-up fashion sneakers with white soles"
}
]}
"""