Spaces:
Sleeping
Sleeping
| 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 | |
| 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) | |