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())
|