Spaces:
Running
Running
File size: 3,419 Bytes
74b4285 c396ba7 74b4285 c396ba7 74b4285 |
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 |
import asyncio
from aiohttp import web, WSMsgType
from google import genai
from google.genai import types
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
import os
def load_pdf_chunks(pdf_path, max_chunk_length=500):
reader = PdfReader(pdf_path)
full_text = ""
for page in reader.pages:
text = page.extract_text()
if text:
full_text += text + "\n"
# Split into smaller chunks
chunks = [full_text[i:i + max_chunk_length] for i in range(0, len(full_text), max_chunk_length)]
return chunks
PDF_PATH = "resume.pdf"
chunks = load_pdf_chunks(PDF_PATH)
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
chunk_embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
dimension = chunk_embeddings.shape[1]
index = faiss.IndexFlatL2(dimension)
index.add(chunk_embeddings)
def retrieve_context(query, top_k=7):
query_embedding = embedding_model.encode([query], convert_to_numpy=True)
distances, indices = index.search(query_embedding, top_k)
return [chunks[i] for i in indices[0]]
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
if not GEMINI_API_KEY:
raise ValueError("GEMINI_API_KEY environment variable not set!")
# Initialize Gemini client
client = genai.Client(api_key=GEMINI_API_KEY)
model = "gemini-2.0-flash-live-001"
config = types.LiveConnectConfig(
response_modalities=["TEXT"],
system_instruction=types.Content(
parts=[
types.Part.from_text(
text="You are Soumyajit's AI Assistant which is chatbot on Soumyajit's profile website, you have to act like a professional AI Assistant, and answer precisely in maximum 2-3 lines."
)
],
role="user"
)
)
# WebSocket handler
async def ws_handler(request):
ws = web.WebSocketResponse()
await ws.prepare(request)
async with client.aio.live.connect(model=model, config=config) as session:
async for msg in ws:
if msg.type == WSMsgType.TEXT:
user_query = msg.data
# 🔍 RAG Retrieval
retrieved_docs = retrieve_context(user_query)
context_text = "\n".join(retrieved_docs)
# Combine into final prompt
final_prompt = f"Context:\n{context_text}\n\nQuestion:\n{user_query}"
await session.send_client_content(
turns={"role": "user", "parts": [{"text": final_prompt}]},
turn_complete=True
)
full_response = ""
async for response in session.receive():
if response.text:
full_response += response.text
await ws.send_str(full_response)
elif msg.type == WSMsgType.ERROR:
print('WebSocket connection closed with exception:', ws.exception())
return ws
# Health check
async def health_check(request):
return web.Response(text="OK")
# Main app setup
app = web.Application()
app.router.add_get("/", health_check)
app.router.add_get("/ws", ws_handler) # WebSocket endpoint
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860)) # Hugging Face uses 7860
web.run_app(app, port=port, host="0.0.0.0") |