Spaces:
Running
Running
File size: 5,040 Bytes
d2aa2a5 5595a23 d2aa2a5 adb38e8 d2aa2a5 adb38e8 d2aa2a5 adb38e8 d2aa2a5 adb38e8 d2aa2a5 adb38e8 d2aa2a5 adb38e8 | 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 | from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from huggingface_hub import InferenceClient
import os, logging, time, threading
from loader import Loader
from chunker import Chunker
from embedder import Embedder
from vector import VectorStorage
from retriever import Retriever
app = FastAPI()
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
MODELS = [
"Qwen/Qwen2.5-72B-Instruct",
"meta-llama/Llama-3.2-3B-Instruct",
"deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
"mistralai/Mistral-7B-Instruct-v0.3",
"HuggingFaceH4/zephyr-7b-beta",
]
SESSION_TIMEOUT = 3 * 60 * 60
sessions: dict = {}
def auto_cleanup():
while True:
time.sleep(SESSION_TIMEOUT)
current_time = time.time()
expired = [
sid for sid, data in sessions.items()
if current_time - data.get('created_at', current_time) > SESSION_TIMEOUT
]
for sid in expired:
del sessions[sid]
if expired:
logger.info(f"Auto-cleaned {len(expired)} expired sessions to free RAM.")
threading.Thread(target=auto_cleanup, daemon=True).start()
# --- Global RAG Components ---
text = Loader("portfolio.pdf").load()
chunks = Chunker().chunker(text)
embedder = Embedder()
vectors = embedder.embed(chunks)
store = VectorStorage(dimension=len(vectors[0]))
store.add(vectors, chunks)
class ChatRequest(BaseModel):
session_id: str
message: str
@app.post("/")
def chat(req: ChatRequest):
if req.session_id not in sessions:
sessions[req.session_id] = {"history": [], "created_at": time.time()}
session = sessions[req.session_id]
retriever = Retriever(store, embedder, k=3)
context_chunks = retriever.retrieve(req.message)
if not context_chunks:
return {"response": "I only answer questions about Aarav and his work."}
context_text = "\n\n".join(context_chunks)
system_prompt = (
"You are Aarav's AI assistant.\n"
"Your name is Zooba\n"
"Your job is to answer questions about Aarav Kumar Ranjan, his projects, skills, and interests using the provided context.\n"
"Rules:\n"
"- Only answer using the given context. Do not make up information.\n"
"- If the answer is not in the context, say: I only answer questions about Aarav and his work.\n"
"- Keep answers clear, simple, and confident.\n"
"- Do not use complex jargon unless necessary.\n"
"- Prefer explaining things in a way a beginner can understand.\n"
"Style:\n"
"- Speak in a calm, intelligent, and slightly friendly tone.\n"
"- Be concise but informative.\n"
"- When explaining projects, include:\n"
" • what it does\n"
" • how it works (simple explanation)\n"
" • why it is useful\n"
"Do not generate fake achievements, skills, or experiences.\n"
"Do not pretend to be Aarav himself.\n"
"If asked about projects, mention their names clearly.\n"
"Make Aarav appear as a thoughtful, skilled, and curious machine learning enthusiast who focuses on understanding and building real systems.\n"
)
messages = [{"role": "system", "content": system_prompt}]
recent_history = session["history"][-10:]
messages.extend(recent_history)
messages.append({"role": "user", "content": f"Context:\n{context_text}\n\nQuestion: {req.message}"})
full_response = ""
def token_stream():
for model in MODELS:
try:
client = InferenceClient(model, token=os.environ["HF_TOKEN"])
logger.info(f"Streaming with: {model}")
success = False
for token in client.chat_completion(messages, max_tokens=512, stream=True):
text = token.choices[0].delta.content
if text:
success = True
yield f"data: {text}\n\n"
yield "data: [DONE]\n\n"
return
except Exception as e:
if success:
yield "data: [DONE]\n\n"
return
logger.warning(f"Streaming failed for {model}: {e}")
continue
yield "data: Sorry, all models are currently unavailable. Try again later.\n\n"
yield "data: [DONE]\n\n"
return StreamingResponse(token_stream(), media_type="text/event-stream")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7600)
|