hank1229's picture
Update llm.py
f91e2fa verified
import json
from transformers import pipeline
class StyleSavvy:
def __init__(
self,
model_name: str = "google/flan-t5-large",
device: int = -1, # -1 = CPU, or GPU index
max_length: int = 150,
):
self.pipe = pipeline(
"text2text-generation",
model=model_name,
tokenizer=model_name,
device=device,
)
self.max_length = max_length
def advise(self, items, body_type, face_shape, occasion):
# 如果 items 是 JSON 字串,先 parse 回 Python list
items = json.loads(items) if isinstance(items, str) else items
prompt = (
f"The user is {body_type}-shaped with a {face_shape} face, "
f"attending a {occasion}. They are wearing: "
+ ", ".join(i["label"] for i in items)
+ ".\n\nPlease list 5 concise style tips as bullet points:"
)
result = self.pipe(
prompt,
max_length=self.max_length,
num_beams=4,
early_stopping=True,
do_sample=False
)[0]["generated_text"].strip()
return result
# import torch
# # models/llm.py
# # models/llm.py
# import os
# from typing import List
# from transformers import pipeline, Pipeline
# # Force CPU modes (avoid any MPS/CUDA issues on macOS)
# os.environ["CUDA_VISIBLE_DEVICES"] = ""
# os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "0"
# class StyleSavvy:
# def __init__(
# self,
# model_name: str = "openlm-research/open_llama_3b_v2",
# device: int = -1, # -1 = CPU
# max_new_tokens: int = 100,
# temperature: float = 0.7,
# top_p: float = 0.9,
# ):
# """
# Uses OpenLLaMA-3B-v2 (≈3B params) for fast, local inference.
# """
# # Setup a causal text-generation pipeline
# self.pipe: Pipeline = pipeline(
# "text-generation",
# model=model_name,
# tokenizer=model_name,
# device=device,
# )
# # GPT‐style models need a pad token to avoid warnings
# if self.pipe.tokenizer.pad_token_id is None:
# self.pipe.tokenizer.pad_token = self.pipe.tokenizer.eos_token
# self.max_new_tokens = max_new_tokens
# self.temperature = temperature
# self.top_p = top_p
# def advise(
# self,
# items: List[str],
# body_type: str,
# face_shape: str,
# occasion: str
# ) -> List[str]:
# """
# Builds a strict instruction prompt and returns exactly five "- " bullets.
# """
# labels = ", ".join(items) if items else "an outfit"
# prompt = (
# "You are a professional fashion consultant.\n"
# f"The user is {body_type}-shaped with a {face_shape} face, attending {occasion}.\n"
# f"They are wearing: {labels}.\n\n"
# "Please provide exactly five concise style tips, each on its own line, "
# "and starting with \"- \". No extra text."
# )
# # Generate
# output = self.pipe(
# prompt,
# max_new_tokens=self.max_new_tokens,
# do_sample=True,
# temperature=self.temperature,
# top_p=self.top_p,
# return_full_text=False,
# )[0]["generated_text"]
# # Extract bullets
# tips = [ln.strip() for ln in output.splitlines() if ln.strip().startswith("- ")]
# # Fallback: split on sentences if fewer than 5 bullets
# if len(tips) < 5:
# candidates = [s.strip() for s in output.replace("\n"," ").split(".") if s.strip()]
# tips = [f"- {candidates[i]}" for i in range(min(5, len(candidates)))]
# return tips[:5]