Spaces:
Runtime error
Runtime error
File size: 4,988 Bytes
89b6166 |
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 |
"""
Standalone RAG service for Nebius deployment
Replaces Modal-specific code with standard Python/Flask
"""
from flask import Flask, request, jsonify
from flask_cors import CORS
import os
import time
from pathlib import Path
# Import your RAG components (adapted for Nebius)
from langchain_community.embeddings import HuggingFaceEmbeddings
from vllm import LLM, SamplingParams
from langchain.schema import Document
import chromadb
app = Flask(__name__)
CORS(app)
# Configuration
LLM_MODEL = os.getenv("LLM_MODEL", "microsoft/Phi-3-mini-4k-instruct")
EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5")
CHROMA_PERSIST_DIR = os.getenv("CHROMA_PERSIST_DIR", "/data/chroma")
DOCUMENTS_DIR = os.getenv("DOCUMENTS_DIR", "/data/documents")
# Global variables for models (loaded at startup)
embeddings = None
llm_engine = None
sampling_params = None
collection = None
def initialize_models():
"""Initialize models at startup"""
global embeddings, llm_engine, sampling_params, collection
print("🚀 Initializing RAG service...")
# Initialize embeddings
print(" Loading embeddings model...")
embeddings = HuggingFaceEmbeddings(
model_name=EMBEDDING_MODEL,
model_kwargs={'device': 'cuda'},
encode_kwargs={'normalize_embeddings': True}
)
# Initialize LLM
print(" Loading LLM...")
llm_engine = LLM(
model=LLM_MODEL,
dtype="float16",
gpu_memory_utilization=0.85,
max_model_len=4096,
trust_remote_code=True,
enforce_eager=True
)
sampling_params = SamplingParams(
temperature=0.7,
max_tokens=1536,
top_p=0.9,
stop=["\n\n\n", "Question:", "Context:", "<|end|>"]
)
# Initialize ChromaDB
print(" Connecting to ChromaDB...")
chroma_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR)
collection = chroma_client.get_or_create_collection("product_design")
print("✅ RAG service ready!")
# Initialize on startup
initialize_models()
@app.route('/health', methods=['GET'])
def health():
"""Health check endpoint"""
return jsonify({
"status": "healthy",
"llm_model": LLM_MODEL,
"embedding_model": EMBEDDING_MODEL
}), 200
@app.route('/query', methods=['POST'])
def query():
"""Query the RAG system"""
data = request.json
question = data.get('question', '').strip()
if not question:
return jsonify({"error": "Question is required"}), 400
start_time = time.time()
try:
# Generate query embedding
query_embedding = embeddings.embed_query(question)
# Retrieve from ChromaDB
results = collection.query(
query_embeddings=[query_embedding],
n_results=5
)
retrieval_time = time.time() - start_time
if not results['documents'] or len(results['documents'][0]) == 0:
return jsonify({
"question": question,
"answer": "No relevant information found in the product design document.",
"retrieval_time": retrieval_time,
"generation_time": 0,
"sources": []
})
# Build context
context = "\n\n".join(results['documents'][0])
# Generate prompt
prompt = f"""<|system|>
You are a helpful AI assistant that answers questions about the TokyoDrive Insurance product design document.
Provide comprehensive, detailed answers with specific information from the document.
Structure your answer clearly with specific numbers, percentages, and data points.
Be thorough and cite specific details from the context. If information is not available, say so clearly.<|end|>
<|user|>
Context from Product Design Document:
{context}
Question:
{question}<|end|>
<|assistant|>"""
# Generate answer
gen_start = time.time()
outputs = llm_engine.generate(prompts=[prompt], sampling_params=sampling_params)
answer = outputs[0].outputs[0].text.strip()
generation_time = time.time() - gen_start
# Prepare sources
sources = []
for i, doc in enumerate(results['documents'][0]):
sources.append({
"content": doc[:500],
"metadata": results.get('metadatas', [[{}]])[0][i] if results.get('metadatas') else {}
})
return jsonify({
"question": question,
"answer": answer,
"retrieval_time": retrieval_time,
"generation_time": generation_time,
"sources": sources,
"success": True
})
except Exception as e:
return jsonify({
"success": False,
"error": str(e)
}), 500
if __name__ == '__main__':
port = int(os.getenv("PORT", 8000))
app.run(host='0.0.0.0', port=port, debug=False)
|