Update main.py

main.py

@@ -58,16 +58,15 @@ from typing import List
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from sentence_transformers import SentenceTransformer
 from pypdf import PdfReader
-import torch
-import os
+import torch, os
 
-# ======================================================
-# FastAPI App
-# ======================================================
+# ======================================
+# FastAPI Base
+# ======================================
 app = FastAPI(
     title="Harshal AI Backend",
     version="1.0.0",
-    description="Human-like AI for Harshal Portfolio"
+    description="Human-like assistant bound to Harshal's real resume facts."
 )
 
 app.add_middleware(
@@ -77,123 +76,109 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# ======================================================
-# 1) LOAD MAIN MODEL — Qwen2.5 1.5B (CPU Friendly)
-# ======================================================
+# ======================================
+# MODEL (Qwen2.5 1.5B)
+# ======================================
 MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
 
-print(f"🚀 Loading LLM: {MODEL_NAME}")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-
 llm = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    dtype=torch.float32,      # correct argument, CPU-safe
+    dtype=torch.float32,
+    device_map="cpu"
 )
 llm.eval()
 
-print("✅ Qwen Loaded Successfully")
-
-# ======================================================
-# 2) LOAD RESUME + RAG
-# ======================================================
-EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-embedder = SentenceTransformer(EMBED_MODEL)
-
-RESUME_PATH = "resume.pdf"
-resume_rag = None
-
-
-def chunk_text(text, max_chars=450, overlap=80):
+# ======================================
+# EMBEDDING MODEL
+# ======================================
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+resume_data = None
+RESUME_FILE = "resume.pdf"
+
+# ======================================
+# UTIL: Chunk Resume
+# ======================================
+def chunk_text(text, size=450, overlap=80):
     text = " ".join(text.split())
-    chunks = []
-    start = 0
+    out, start = [], 0
+
     while start < len(text):
-        end = start + max_chars
-        chunks.append(text[start:end])
+        end = start + size
+        out.append(text[start:end])
         start = end - overlap
-    return chunks
-
 
-def build_resume_rag():
-    """Read resume.pdf, chunk it, embed it."""
-    global resume_rag
+    return out
 
-    if not os.path.exists(RESUME_PATH):
-        print("⚠ resume.pdf not found — RAG disabled")
+# ======================================
+# BUILD RAG
+# ======================================
+def load_resume():
+    global resume_data
+    if not os.path.exists(RESUME_FILE):
+        print("❌ resume.pdf not found")
         return
 
-    try:
-        reader = PdfReader(RESUME_PATH)
-        full = ""
-        for p in reader.pages:
-            full += p.extract_text() or ""
+    reader = PdfReader(RESUME_FILE)
+    text = ""
+    for pg in reader.pages:
+        text += pg.extract_text() or ""
 
-        chunks = chunk_text(full)
-        embs = embedder.encode(
-            chunks,
-            convert_to_tensor=True,
-            normalize_embeddings=True
-        )
-        resume_rag = {"chunks": chunks, "embs": embs}
-        print(f"✅ Resume RAG ready ({len(chunks)} chunks)")
+    chunks = chunk_text(text)
+    embs = embedder.encode(chunks, convert_to_tensor=True, normalize_embeddings=True)
+    resume_data = {"chunks": chunks, "embs": embs}
 
-    except Exception as e:
-        print("❌ RAG build error:", e)
-        resume_rag = None
+    print(f"✅ Resume RAG loaded with {len(chunks)} chunks")
 
 
-def get_rag_context(query: str, top_k=3):
-    if resume_rag is None:
+def get_resume_context(query, top_k=3):
+    if resume_data is None:
         return ""
 
     q = embedder.encode([query], convert_to_tensor=True, normalize_embeddings=True)[0]
-    sims = torch.nn.functional.cosine_similarity(q.unsqueeze(0), resume_rag["embs"])
-    idx = torch.topk(sims, k=min(top_k, len(sims))).indices
+    sims = torch.nn.functional.cosine_similarity(q.unsqueeze(0), resume_data["embs"])
+    top = torch.topk(sims, k=top_k)
 
-    return "\n\n".join(resume_rag["chunks"][i] for i in idx)
+    return "\n\n".join(resume_data["chunks"][i] for i in top.indices)
 
 
-build_resume_rag()
+load_resume()
 
-# ======================================================
-# 3) API BODY DEFINITIONS — CLEAN SWAGGER
-# ======================================================
-class Message(BaseModel):
+# ======================================
+# SCHEMAS
+# ======================================
+class Msg(BaseModel):
     role: str
     content: str
 
+class ChatReq(BaseModel):
+    messages: List[Msg]
 
-class ChatRequest(BaseModel):
-    messages: List[Message]
-
-
-# ======================================================
-# 4) CHAT ROUTE — HUMAN-LIKE HARSHAAL AI
-# ======================================================
+# ======================================
+# CHAT
+# ======================================
 @app.post("/chat")
-def chat(req: ChatRequest):
-
-    if not req.messages:
-        return {"reply": "Hey, ask me something 🙂"}
+def chat(req: ChatReq):
 
     user_msg = req.messages[-1].content.strip()
+    resume_context = get_resume_context(user_msg)
 
-    # -------- Retrieve relevant resume context --------
-    rag_text = get_rag_context(user_msg)
-
-    # -------- Persona --------
+    # STRICT anti-hallucination persona
     persona = f"""
 You are Harshal Sonawane, a real human software engineer from Pune.
-Speak naturally, casual but professional.
-Never say you're an AI or model.
-Use resume context when relevant.
-Keep replies short (2–4 sentences).
 
-Resume knowledge:
-{rag_text}
+VERY IMPORTANT RULES:
+- NEVER say you are an AI, model, assistant, chatbot.
+- ONLY answer using information available in the resume context.
+- If the resume does NOT contain the answer → say: "I haven't mentioned this in my resume."
+- Do NOT invent employment history, projects, companies, or skills.
+- Use a natural, short, human tone (2–4 sentences).
+- Stay consistent with resume facts.
+
+Resume facts (DO NOT IGNORE):
+{resume_context}
 """.strip()
 
-    # -------- Model Chat Format --------
     messages = [
         {"role": "system", "content": persona},
         {"role": "user", "content": user_msg}
@@ -202,27 +187,26 @@ Resume knowledge:
     ids = tokenizer.apply_chat_template(
         messages,
         return_tensors="pt",
-        add_generation_prompt=True,
+        add_generation_prompt=True
+    ).to(llm.device)
+
+    out = llm.generate(
+        ids,
+        max_new_tokens=160,
+        temperature=0.45,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     )
 
-    with torch.no_grad():
-        out = llm.generate(
-            ids,
-            max_new_tokens=150,
-            temperature=0.45,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            do_sample=True
-        )
-
     reply = tokenizer.decode(out[0][ids.shape[-1]:], skip_special_tokens=True).strip()
 
     return {"reply": reply}
 
-
-# ======================================================
-# 5) HEALTH CHECK
-# ======================================================
+# ======================================
+# HEALTH
+# ======================================
 @app.get("/")
-def root():
-    return {"status": "Harshal AI backend is running (Qwen2.5 + RAG) 🎯"}
+def health():
+    return {"status": "Harshal AI backend running with Qwen 1.5B + strict RAG 🎯"}