62 lines
3.0 KiB
Python
62 lines
3.0 KiB
Python
|
from langchain.output_parsers import ResponseSchema, StructuredOutputParser
|
|||
|
|
from langchain_community.chat_models import ChatTongyi
|
|||
|
|
from langchain_core.prompts import PromptTemplate
|
|||
|
|
|
|||
|
|
from app.schemas.project_info_extraction import ProjectInfoExtractionModel
|
|||
|
|
|
|||
|
|
style = ['NEW_CHINESE', 'COUNTRY_STYLE', 'FUTURISM', 'MINIMALISM', 'LOLITA', 'Y2K', 'BUSINESS', 'MERLAD',
|
|||
|
|
'OUTDOOR_FUNCTIONAL', 'ROCK', 'DOPAMINE', 'GOTHIC', 'POST_APOCALYPTIC', 'ROMANTIC', 'WABI_SABI']
|
|||
|
|
position = ['Overall', 'Tops', 'Bottoms', 'Outwear', 'Blouse', 'Dress', 'Trousers', 'Skirt']
|
|||
|
|
gender = ['Female', 'Male']
|
|||
|
|
age_group = ['Adult', 'Child']
|
|||
|
|
process = ['SERIES_DESIGN', 'SINGLE_DESIGN']
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ProjectInfoExtraction:
|
|||
|
|
def __init__(self, request_data):
|
|||
|
|
# llm generate brand info init
|
|||
|
|
self.model = ChatTongyi(model="qwen2.5-14b-instruct", api_key="sk-7658298c6b99443c98184a5e634fe6ab")
|
|||
|
|
|
|||
|
|
self.response_schemas = [
|
|||
|
|
ResponseSchema(name="project_name", description="项目的名称."),
|
|||
|
|
ResponseSchema(name="process", description="项目的类型，单品还是系列."),
|
|||
|
|
ResponseSchema(name="ageGroup", description="项目设计服装的受众对象."),
|
|||
|
|
ResponseSchema(name="gender", description="项目设计服装的受众性别."),
|
|||
|
|
ResponseSchema(name="position", description="项目单品设计的部位."),
|
|||
|
|
ResponseSchema(name="style", description="项目的设计风格.")
|
|||
|
|
]
|
|||
|
|
self.output_parser = StructuredOutputParser.from_response_schemas(self.response_schemas)
|
|||
|
|
self.format_instructions = self.output_parser.get_format_instructions()
|
|||
|
|
self.prompt = PromptTemplate(
|
|||
|
|
template="你是一个时装品牌的设计师助理。根据用户输入提取出"
|
|||
|
|
"[project_name] : 项目的名称,"
|
|||
|
|
f"[process] : 项目的类型，从{process}选择."
|
|||
|
|
f"[ageGroup] : 服装的受众，从{age_group}选择."
|
|||
|
|
f"[gender] : 服装的适用性别，从{gender}选择"
|
|||
|
|
f"[position] : single_design的部位，如果[process]是SINGLE_DESIGN,从{position}中选择，如果[process]是SERIES_DESIGN，这项为空"
|
|||
|
|
f"[style] : 设计的风格，从{style}中选择"
|
|||
|
|
".\n{format_instructions}\n{question}",
|
|||
|
|
input_variables=["question"],
|
|||
|
|
partial_variables={"format_instructions": self.format_instructions}
|
|||
|
|
)
|
|||
|
|
self._input = self.prompt.format_prompt(question=request_data.prompt)
|
|||
|
|
|
|||
|
|
self.result_data = {}
|
|||
|
|
|
|||
|
|
def get_result(self):
|
|||
|
|
self.llm_extraction_project_info()
|
|||
|
|
return self.result_data
|
|||
|
|
|
|||
|
|
def llm_extraction_project_info(self):
|
|||
|
|
output = self.model(self._input.to_messages())
|
|||
|
|
project_info = self.output_parser.parse(output.content)
|
|||
|
|
self.result_data = project_info
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
request_data = ProjectInfoExtractionModel(
|
|||
|
|
prompt="海边派对主题的衬衫设计"
|
|||
|
|
)
|
|||
|
|
service = ProjectInfoExtraction(request_data)
|
|||
|
|
print(service.get_result())
|