Update main.py

main.py

@@ -58,10 +58,11 @@ from typing import List
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from pypdf import PdfReader
-import torch, os, re
+import torch, os
 
 app = FastAPI(title="Harshal AI Backend", version="1.0.0")
 
+# CORS (Next.js frontend)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -69,163 +70,183 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# -------------------------------------------------------
-# 1. MODEL (CPU SAFE)
-# -------------------------------------------------------
-MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-llm = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
-llm.eval()
-
-# -------------------------------------------------------
-# 2. SMART RAG (Semantic Sections)
-# -------------------------------------------------------
-embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+# ============================================================
+# 1) LOAD MAIN MODEL (Phi-3 Mini — good balance of quality/speed)
+# ============================================================
 
-RESUME = "resume.pdf"
-SECTIONS = {}  # { section_name: text }
+MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
 
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+llm = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    dtype=torch.float32,        # instead of torch_dtype
+)
+llm.eval()
 
-def extract_sections(text):
-    """
-    Extracts logical resume sections using headings.
-    """
-    parts = re.split(r"(PROFESSIONAL EXPERIENCE|PROJECTS|SKILLS|EDUCATION|CERTIFICATION)", text)
-    cleaned = {}
+# ============================================================
+# 2) LOAD EMBEDDINGS + BUILD RAG FROM resume.pdf
+# ============================================================
 
-    current = None
-    for p in parts:
-        p = p.strip()
-        if p in ["PROFESSIONAL EXPERIENCE", "PROJECTS", "SKILLS", "EDUCATION", "CERTIFICATION"]:
-            current = p
-            cleaned[current] = ""
-        elif current:
-            cleaned[current] += p + "\n"
+EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+embedder = SentenceTransformer(EMBED_MODEL)
 
-    return cleaned
+RESUME_FILE = "resume.pdf"
+resume_rag = None
 
 
-def build_rag_sections():
-    global SECTIONS
+def chunk_text(text, max_chars=450, overlap=80):
+    """Simple overlapping chunks."""
+    text = " ".join(text.split())
+    chunks, start = [], 0
 
-    if not os.path.exists(RESUME):
-        print("❌ resume.pdf not found")
-        return
+    while start < len(text):
+        end = start + max_chars
+        chunks.append(text[start:end])
+        start = end - overlap
 
-    reader = PdfReader(RESUME)
-    text = ""
+    return chunks
 
-    for p in reader.pages:
-        text += p.extract_text() or ""
 
-    # Extract using bold headings
-    SECTIONS = extract_sections(text)
-    print("📌 Extracted Sections:", SECTIONS.keys())
+def build_rag():
+    """Reads resume.pdf → chunks → embeddings."""
+    global resume_rag
 
+    if not os.path.exists(RESUME_FILE):
+        print("⚠ resume.pdf NOT FOUND — RAG disabled.")
+        return
 
-build_rag_sections()
+    reader = PdfReader(RESUME_FILE)
+    full_text = ""
 
-# Embed section titles
-SECTION_EMBS = {}
-for k, v in SECTIONS.items():
-    SECTION_EMBS[k] = embedder.encode([v], convert_to_tensor=True)
+    for page in reader.pages:
+        full_text += page.extract_text() or ""
 
+    chunks = chunk_text(full_text)
 
-def retrieve_section(query):
-    """
-    Best matching resume section.
-    """
-    q = embedder.encode([query], convert_to_tensor=True)
-    best = None
-    best_score = -1
+    embeddings = embedder.encode(
+        chunks, convert_to_tensor=True, normalize_embeddings=True
+    )
 
-    for name, emb in SECTION_EMBS.items():
-        sim = torch.nn.functional.cosine_similarity(q, emb)[0].item()
-        if sim > best_score:
-            best_score = sim
-            best = name
+    resume_rag = {"chunks": chunks, "embs": embeddings}
+    print("✅ Resume RAG built with", len(chunks), "chunks")
 
-    return best
 
+build_rag()
 
-# SPECIAL RULES
-def special_logic(query):
-    q = query.lower()
 
-    if any(w in q for w in ["current", "working now", "present job", "company"]):
-        return "PROFESSIONAL EXPERIENCE"
+def retrieve_rag(query, top_k=3):
+    """Find most relevant resume chunks."""
+    if resume_rag is None:
+        return ""
 
-    if any(w in q for w in ["project", "projects", "personal project"]):
-        return "PROJECTS"
+    q = embedder.encode([query], convert_to_tensor=True, normalize_embeddings=True)[0]
+    sims = torch.nn.functional.cosine_similarity(q.unsqueeze(0), resume_rag["embs"])
+    top = torch.topk(sims, k=min(top_k, len(sims)))
 
-    if any(w in q for w in ["skill", "skills"]):
-        return "SKILLS"
+    return "\n\n".join(resume_rag["chunks"][i] for i in top.indices)
 
-    return None
 
+# ============================================================
+# 3) REQUEST MODELS
+# ============================================================
 
-# -------------------------------------------------------
-# 3. CHAT API
-# -------------------------------------------------------
-class Msg(BaseModel):
+class Message(BaseModel):
     role: str
     content: str
 
 class ChatReq(BaseModel):
-    messages: List[Msg]
+    messages: List[Message]
+
 
+# ============================================================
+# 4) SMART MODE CHAT ROUTE (Resume + Fallback)
+# ============================================================
 
 @app.post("/chat")
 def chat(req: ChatReq):
-
     user_msg = req.messages[-1].content.strip()
 
-    # Special rule override
-    sec = special_logic(user_msg)
-    if sec is None:
-        sec = retrieve_section(user_msg)
-
-    resume_text = SECTIONS.get(sec, "")
-
-    persona = f"""
-You are Harshal Sonawane, a real human software engineer from Pune.
-
-Rules:
-- Answer like a friendly human (2–4 sentences).
-- ONLY use information from the resume section below.
-- If the resume does NOT mention the answer:
-    say: "This is not mentioned in my resume, but based on my experience..." 
-    and reply realistically WITHOUT inventing fake jobs.
-
-Resume relevant section ({sec}):
-{resume_text}
+    # RAG: take from resume
+    rag_context = retrieve_rag(user_msg)
+
+    # Detect if we found resume info
+    has_relevant_rag = rag_context and len(rag_context) > 40
+
+    # SMART RULE:
+    #   If resume has relevant info → answer ONLY from resume
+    #   If resume has nothing → fallback using general persona
+    if has_relevant_rag:
+        fallback_note = ""
+        resume_instruction = """
+Use ONLY the details provided in the resume context.
+Do NOT invent any new companies, dates, or roles.
+Stay strictly consistent with:
+- Current company: Börm Bruckmeier Infotech India Pvt. Ltd.
+- Past experience: SarvaTech Consultants (Internship)
+- Personal projects listed in resume.
+"""
+    else:
+        fallback_note = """
+The resume does NOT mention this. 
+Provide a short helpful answer based on Harshal’s background, tone, and skills — 
+BUT do NOT invent companies, dates, or job titles not present in the resume.
+"""
+        resume_instruction = ""
+
+    # System persona
+    system_prompt = f"""
+You are Harshal Sonawane — a real human software engineer from Pune, India.
+
+Tone:
+- Natural, confident, friendly.
+- 2–4 sentences max.
+- Never say you are an AI.
+
+Your resume (RAG context):
+{rag_context}
+
+Instructions:
+{resume_instruction}
+
+Fallback rule:
+{fallback_note}
 """.strip()
 
-    messages = [
-        {"role": "system", "content": persona},
+    # Chat template
+    msgs = [
+        {"role": "system", "content": system_prompt},
         {"role": "user", "content": user_msg},
     ]
 
-    ids = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
-
-    out = llm.generate(
-        ids,
-        max_new_tokens=150,
-        temperature=0.55,
-        top_p=0.9,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
+    input_ids = tokenizer.apply_chat_template(
+        msgs,
+        return_tensors="pt",
+        add_generation_prompt=True,
     )
 
-    gen = out[0][ids.shape[-1]:]
+    with torch.no_grad():
+        output_ids = llm.generate(
+            input_ids,
+            max_new_tokens=160,
+            temperature=0.55,
+            top_p=0.9,
+            repetition_penalty=1.06,
+            do_sample=True,
+        )
+
+    gen = output_ids[0][input_ids.shape[-1]:]
     reply = tokenizer.decode(gen, skip_special_tokens=True).strip()
 
-    if resume_text.strip() == "":
-        reply = "This is not mentioned in my resume, but based on my experience, " + reply
-
     return {"reply": reply}
 
 
+# ============================================================
+# 5) HEALTH CHECK
+# ============================================================
+
 @app.get("/")
 def health():
-    return {"status": "Harshal AI running with smart RAG 🎯"} 
+    return {
+        "status": "Harshal AI backend running (SMART MODE)",
+        "model": MODEL_NAME
+    }