Spaces:
Running on Zero
Running on Zero
File size: 3,905 Bytes
a65585a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import torch
from datasets import load_dataset
from transformers import (
AutoModel,
AutoModelForSequenceClassification,
AutoProcessor,
Qwen3VLForConditionalGeneration,
)
from pinecone import Pinecone
from redisvl.index import SearchIndex
from redisvl.schema import IndexSchema
from config import (
DEVICE,
EMBED_MODEL_PATH,
EMBED_COMMIT_HASH,
RERANK_MODEL_PATH,
RERANK_COMMIT_HASH,
GENERATION_MODEL_ID,
PINECONE_API_KEY,
PINECONE_INDEX_NAME,
REDIS_URL,
REDIS_DBNAME
)
# ============================================================================
# Initialize External Services
# ============================================================================
print("[INFO] Initializing Pinecone...")
pc = Pinecone(api_key=PINECONE_API_KEY) if PINECONE_API_KEY else None
pinecone_index = pc.Index(PINECONE_INDEX_NAME) if pc else None
print("[INFO] Initializing RedisVL SearchIndex...")
try:
schema_dict = {
"index": {
"name": REDIS_DBNAME,
"prefix": "cache"
},
"fields": [
{"name": "prompt", "type": "text"},
{"name": "response", "type": "text"},
{"name": "vector", "type": "vector", "attrs": {"dims": 2048, "distance_metric": "cosine", "algorithm": "flat", "datatype": "float32"}}
]
}
schema = IndexSchema.from_dict(schema_dict)
redis_cache = SearchIndex(schema, redis_url=REDIS_URL)
redis_cache.create(overwrite=True)
except Exception as e:
print(f"[WARNING] Could not initialize RedisVL SearchIndex: {e}")
redis_cache = None
# ============================================================================
# Load Dataset
# ============================================================================
print("[INFO] Loading dataset...")
dataset = load_dataset(path="mrdbourke/recipe-synthetic-images-10k")
print(f"[INFO] Dataset loaded with {len(dataset['train'])} samples")
# ============================================================================
# Load Models
# ============================================================================
modality_to_tokens = {
"image": 2048,
"image_text": 10240,
"text": 8192
}
print(f"[INFO] Loading embedding model from: {EMBED_MODEL_PATH} with commit: {EMBED_COMMIT_HASH}")
embed_model = AutoModel.from_pretrained(
EMBED_MODEL_PATH,
revision=EMBED_COMMIT_HASH,
dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="sdpa",
device_map="auto",
).eval()
embed_modality = "image_text"
embed_processor_kwargs = {
"max_input_tiles": 6,
"use_thumbnail": True,
"p_max_length": modality_to_tokens[embed_modality]
}
embed_processor = AutoProcessor.from_pretrained(
EMBED_MODEL_PATH,
revision=EMBED_COMMIT_HASH,
trust_remote_code=True,
**embed_processor_kwargs
)
print(f"[INFO] Embedding model loaded!")
print(f"[INFO] Loading rerank model from: {RERANK_MODEL_PATH} with commit: {RERANK_COMMIT_HASH}")
rerank_model = AutoModelForSequenceClassification.from_pretrained(
RERANK_MODEL_PATH,
revision=RERANK_COMMIT_HASH,
dtype=torch.bfloat16,
trust_remote_code=True,
attn_implementation="sdpa",
device_map="auto",
).eval()
rerank_modality = "image_text"
rereank_processor_kwargs = {
"max_input_tiles": 6,
"use_thumbnail": True,
"rerank_max_length": modality_to_tokens[rerank_modality]
}
rerank_processor = AutoProcessor.from_pretrained(
RERANK_MODEL_PATH,
revision=RERANK_COMMIT_HASH,
trust_remote_code=True,
**rereank_processor_kwargs
)
print(f"[INFO] Rerank model loaded!")
print("[INFO] Loading generation model...")
qwen_model = Qwen3VLForConditionalGeneration.from_pretrained(
GENERATION_MODEL_ID,
dtype="auto",
device_map="auto"
)
qwen_processor = AutoProcessor.from_pretrained(GENERATION_MODEL_ID)
print(f"[INFO] Generation model loaded")
|