Update app.py

app.py

@@ -2,49 +2,41 @@ import gradio as gr
 from transformers import pipeline
 from datasets import load_dataset
 
-# Load dataset and models
-dataset = load_dataset("AntZet/home_decoration_objects_images")
-style_clf = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-advice_gen = pipeline("text2text-generation", model="google/flan-t5-base")
+# 加载数据集
+dataset = load_dataset("AntZet/home_decoration_objects_images", streaming=True)
+dataset = dataset['train'].take(100)  # 只取前100个样本加快加载
 
-def get_recommendation(style_input):
-    # 1. Find matching styles
-    style_examples = [d for d in dataset['train'] if d['style'].lower() == style_input.lower()]
-    
-    if not style_examples:
-        return "Style not found. Try: Industrial, Scandinavian, etc."
-    
-    # 2. Get common objects
-    objects = [obj for ex in style_examples for obj in ex['objects']]
-    top_objects = sorted(set(objects), key=lambda x: -objects.count(x))[:3]
+# 初始化模型
+style_advisor = pipeline("text-generation", model="gpt2")  # 使用较小的gpt2模型
+
+def get_advice(style):
+    # 从数据集中找出匹配风格的例子
+    examples = [ex for ex in dataset if ex['style'].lower() == style.lower()]
     
-    # 3. Generate advice
-    prompt = f"""As an Airbnb host, how to decorate in {style_input} style?
-    Key objects: {top_objects}
-    Provide:
-    - 3 essential elements
-    - 2 budget tips
-    - Colors to use"""
+    if not examples:
+        return "未找到该风格，请尝试：工业风、北欧风等", []
     
-    advice = advice_gen(prompt, max_length=300)[0]['generated_text']
+    # 生成建议
+    prompt = f"如何将Airbnb房间装修成{style}风格？请给出3条具体建议"
+    advice = style_advisor(prompt, max_length=150)[0]['generated_text']
     
-    # 4. Show example images
-    example_images = [ex['image'] for ex in style_examples[:3]]
+    # 获取示例图片
+    example_images = [ex['image'] for ex in examples[:3]]
     
     return advice, example_images
 
-# Simple UI
+# 创建界面
 with gr.Blocks() as demo:
-    gr.Markdown("## 🏡 Airbnb Style Advisor")
-    style_input = gr.Textbox(label="Enter style (e.g. Industrial)")
-    submit_btn = gr.Button("Get Advice")
+    gr.Markdown("## 🏡 Airbnb装修助手")
+    style_input = gr.Textbox(label="输入想要的风格")
+    submit_btn = gr.Button("生成建议")
     
     with gr.Row():
-        advice_output = gr.Textbox(label="Recommendations")
-        gallery = gr.Gallery(label="Style Examples")
+        advice_output = gr.Textbox(label="装修建议")
+        gallery = gr.Gallery(label="风格示例")
     
     submit_btn.click(
-        fn=get_recommendation,
+        fn=get_advice,
         inputs=style_input,
         outputs=[advice_output, gallery]
     )