Spaces:
Sleeping
Sleeping
File size: 7,724 Bytes
dba9e4f e4cbf3a 3820eb5 dba9e4f e4cbf3a dba9e4f e4cbf3a dba9e4f e4cbf3a dba9e4f e4cbf3a dba9e4f e4cbf3a dba9e4f e4cbf3a dba9e4f e4cbf3a 3820eb5 e4cbf3a 3820eb5 e4cbf3a 3820eb5 e4cbf3a 3820eb5 dba9e4f 3820eb5 | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | import os
import gradio as gr
from sentence_transformers import SentenceTransformer
from pathlib import Path
import torch
import numpy as np
from fastapi import FastAPI
import uvicorn
import shutil
# =====================================================
# π§Ή CACHE & STORAGE MANAGEMENT
# =====================================================
# Use temporary directories for caching to prevent disk overuse
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["SENTENCE_TRANSFORMERS_HOME"] = "/tmp"
# Optional: clear any pre-existing cache directories
shutil.rmtree(os.path.expanduser("~/.cache/torch"), ignore_errors=True)
shutil.rmtree(os.path.expanduser("~/.cache/huggingface"), ignore_errors=True)
# =====================================================
# π MODEL SETUP
# =====================================================
# Instead of a heavy local snapshot, load the lightweight MiniLM model directly
MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# =====================================================
# π LOAD MODEL
# =====================================================
try:
print(f"π Loading model: {MODEL_NAME}")
model = SentenceTransformer(MODEL_NAME, device=DEVICE)
print(f"β
Model loaded successfully on {DEVICE.upper()}")
except Exception as e:
raise RuntimeError(
f"β Model loading failed. Ensure internet access or cached files are available.\nError: {e}"
)
# =====================================================
# π§ EMBEDDING FUNCTION (for local + API use)
# =====================================================
def encode_text(text: str):
"""
Returns a summary of the embedding vector (for Gradio interface).
"""
if not text or not text.strip():
return {"error": "Please enter some text."}
try:
embedding = model.encode([text])[0]
preview = np.round(embedding[:10], 4).tolist()
return {
"embedding_shape": len(embedding),
"embedding_preview (first 10 values)": preview,
}
except Exception as e:
return {"error": f"Encoding failed: {str(e)}"}
# =====================================================
# βοΈ API FUNCTION (for remote prediction)
# =====================================================
def api_predict(text: str):
"""
Returns the full embedding vector as a list (for API use).
"""
if not text or not text.strip():
return []
try:
embedding = model.encode([text])[0]
return embedding.tolist()
except Exception as e:
return {"error": str(e)}
# =====================================================
# π¨ GRADIO INTERFACE
# =====================================================
demo = gr.Interface(
fn=encode_text,
inputs=gr.Textbox(label="Enter Text", placeholder="Type a sentence to encode..."),
outputs=gr.JSON(label="Embedding Info"),
title="π¬ Sentence Embedding Generator",
description=(
"Encodes sentences into numerical embeddings using "
"<b>MiniLM-L6-v2</b>. "
"Useful for semantic similarity, clustering, or grouping."
),
examples=[
["I love machine learning and natural language processing."],
["The movie was terrible, I wouldn't recommend it."],
["Artificial intelligence is transforming the world."],
],
)
# =====================================================
# π FASTAPI APP + GRADIO MOUNT
# =====================================================
app = FastAPI(title="Sentence Embedding API")
# Add the /predict endpoint
@app.post("/predict")
async def predict_api(data: dict):
text = data.get("text", "")
return api_predict(text)
# Mount Gradio UI at root "/"
app = gr.mount_gradio_app(app, demo, path="/")
# =====================================================
# π₯ LAUNCH SERVER
# =====================================================
if __name__ == "__main__":
print("π Starting API and Gradio UI server...")
uvicorn.run(app, host="0.0.0.0", port=7860)
# import gradio as gr
# from sentence_transformers import SentenceTransformer
# from pathlib import Path
# import torch
# import numpy as np
# # =====================================================
# # π MODEL SETUP
# # =====================================================
# MODEL_DIR = Path(__file__).parent / "models" / "all-MiniLM" / "snapshots" / "c9745ed1d9f207416be6d2e6f8de32d1f16199bf"
# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# # =====================================================
# # π LOAD MODEL
# # =====================================================
# try:
# print(f"π Loading model from: {MODEL_DIR}")
# model = SentenceTransformer(str(MODEL_DIR), device=DEVICE)
# print(f"β
Model loaded successfully on {DEVICE.upper()}")
# except Exception as e:
# raise RuntimeError(
# f"β Model loading failed. Ensure the path '{MODEL_DIR}' "
# f"is correct and contains model files. Original error: {e}"
# )
# # =====================================================
# # π§ EMBEDDING FUNCTION (for local + API use)
# # =====================================================
# def encode_text(text: str):
# """
# Returns the embedding vector of a given text.
# For API use, only returns the vector (list of floats).
# """
# if not text or not text.strip():
# return {"error": "Please enter some text."}
# try:
# embedding = model.encode([text])[0]
# preview = np.round(embedding[:10], 4).tolist()
# return {
# "embedding_shape": len(embedding),
# "embedding_preview (first 10 values)": preview,
# }
# except Exception as e:
# return {"error": f"Encoding failed: {str(e)}"}
# # =====================================================
# # βοΈ API FUNCTION (for remote prediction)
# # =====================================================
# def api_predict(text: str):
# """
# This is the function that will be exposed via /predict API.
# It returns only the embedding vector (as a list).
# """
# if not text or not text.strip():
# return []
# try:
# embedding = model.encode([text])[0]
# return embedding.tolist()
# except Exception as e:
# return {"error": str(e)}
# # =====================================================
# # π¨ GRADIO INTERFACE
# # =====================================================
# demo = gr.Interface(
# fn=encode_text,
# inputs=gr.Textbox(label="Enter Text", placeholder="Type a sentence to encode..."),
# outputs=gr.JSON(label="Embedding Info"),
# title="π¬ Sentence Embedding Generator",
# description=(
# "Encodes sentences into numerical embeddings using "
# "<b>distilbert-base-nli-mean-tokens</b>. "
# "Useful for semantic similarity, clustering, or grouping."
# ),
# examples=[
# ["I love machine learning and natural language processing."],
# ["The movie was terrible, I wouldn't recommend it."],
# ["Artificial intelligence is transforming the world."],
# ],
# )
# # =====================================================
# # π§© ADD API ENDPOINT (/predict)
# # =====================================================
# demo.add_api_route("/predict", api_predict, methods=["POST"])
# # =====================================================
# # π LAUNCH SERVER
# # =====================================================
# if __name__ == "__main__":
# print("π Starting Gradio API server...")
# demo.launch(server_name="0.0.0.0", server_port=7860)
|