Update main.py

main.py

@@ -52,49 +52,101 @@
 
 
 
-
 from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.openapi.utils import get_openapi
 from pydantic import BaseModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from sentence_transformers import SentenceTransformer, util
-from pathlib import Path
-from pypdf import PdfReader
-import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch, os
 
 app = FastAPI()
 
-# Allow your Next.js frontend to call this from anywhere
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# ------------------------------------------
+# 1) USE FASTER + SMARTER MODEL (FREE)
+# ------------------------------------------
+MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
-# -------- LLM --------
-LLM_MODEL = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
 model = AutoModelForCausalLM.from_pretrained(
-    LLM_MODEL,
+    MODEL_NAME,
     torch_dtype=torch.float32,
-    device_map="cpu",
+    device_map="cpu"
 )
 
+# ------------------------------------------
+# 2) RAG LOADING (resume embeddings)
+# ------------------------------------------
+if os.path.exists("vector_store.pt"):
+    vectors = torch.load("vector_store.pt")
+else:
+    vectors = None
+
+def retrieve_relevant_text(query):
+    if vectors is None:
+        return ""
+    # simple cosine search (fast)
+    qs = embed(query)
+    sims = torch.nn.functional.cosine_similarity(qs, vectors['embeddings'])
+    top_idx = sims.argmax().item()
+    return vectors["chunks"][top_idx]
+
+# ------------------------------------------
+# Request format
+# ------------------------------------------
+class ChatRequest(BaseModel):
+    messages: list
+
+# ------------------------------------------
+# 3) MAIN CHAT ROUTE
+# ------------------------------------------
+@app.post("/chat")
+def chat(req: ChatRequest):
+
+    user_msg = req.messages[-1]["content"]
+
+    # RAG context
+    rag_text = retrieve_relevant_text(user_msg)
+
+    # Persona system prompt
+    system_prompt = f"""
+You are Harshal Sonawane — a real human software engineer.
+You speak casually, friendly, and naturally.
+You NEVER say you are an AI or chatbot.
+
+Your goals:
+- Answer clearly and like a human.
+- Use resume knowledge when helpful.
+- Do NOT hallucinate.
+- Keep responses short (2–4 sentences).
+- Stay in character as Harshal.
 
-# -------- Embeddings + Resume store --------
-EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-embedder = SentenceTransformer(EMBED_MODEL)
+Resume context (if relevant):
+{rag_text}
+""".strip()
 
+    full_prompt = f"<|system|>{system_prompt}\n<|user|>{user_msg}\n<|assistant|>"
 
-class ResumeStore:
-    def __init__(self):
-        self.chunks: list[str] = []
-        self.embeddings = None
+    # Generate human-like answer FAST
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+
+    output = model.generate(
+        **inputs,
+        max_new_tokens=120,
+        temperature=0.6,
+        top_p=0.9,
+        repetition_penalty=1.05,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    reply = tokenizer.decode(output[0], skip_special_tokens=True)
+    reply = reply.split("<|assistant|>")[-1].strip()
+
+    return {"reply": reply}
+
+@app.get("/")
+def home():
+    return {"message": "Harshal AI backend running smoothly!"}
 
 
-RESUME = ResumeStore()
 
 
 def load_resume():