from transformers import T5Tokenizer, T5ForConditionalGeneration

# Load T5-large model and tokenizer
model_name = "t5-large"
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)

# Define the input text (structured schema + task)
input_text = """
Schema:
Table: Products
Fields:
    - ProductID (Primary Key, int)
    - ProductName (string)
    - Category (string)
    - Price (float)
    - Tags (string)

Table: Categories
Fields:
    - CategoryID (Primary Key, int)
    - CategoryName (string)

Task: Identify the specific fields required to set up a product similarity search. Return the output in the format:
Example Output:
    [Products.ProductName]
    [Products.Category]
    [Products.Tags]
"""

# Tokenize the input
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)

# Generate output
outputs = model.generate(**inputs, max_length=100, num_beams=4, early_stopping=True)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)

print("Recommended Fields:", result)
print("End of Recommended Fields")