Update app/app.py

app/app.py

@@ -1,96 +1,67 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, Request
 from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from app.policy_vector_db import PolicyVectorDB  # Import your class
+from policy_vector_db import PolicyVectorDB  # Make sure this is your local DB logic
+import chromadb
 
-# --- 1. Initialize the Vector Database and LLM ---
+# Create FastAPI app
+app = FastAPI()
 
-# Load the vector database from /tmp (safest in Docker/HF Spaces)
+# Load the vector database from /tmp (safe for Hugging Face Spaces)
 print("Loading Vector Database...")
 db = PolicyVectorDB(persist_directory="/tmp/policy_vector_db")
 print("Vector Database loaded successfully!")
 
 # Load your quantized model from Hugging Face Hub
-model_id = "Kalpokoch/QuantizedTinyLama"  # Correct spelling assumed
+model_id = "Kalpokoch/QuantizedTinyLama"
 print(f"Loading model: {model_id}...")
+
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-print("Model and tokenizer loaded successfully!")
 
-# Choose dtype depending on device support
-dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+# Quantization config for bitsandbytes
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
 
+# Load quantized model
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=dtype,
-    device_map="auto"
+    device_map="auto",
+    quantization_config=bnb_config
 )
 
-# Create a text-generation pipeline
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=256
-)
+print("Model and tokenizer loaded successfully!")
 
-print("LLM and pipeline loaded successfully!")
 
+# Input schema
+class Query(BaseModel):
+    question: str
 
-# --- 2. FastAPI App Setup ---
-app = FastAPI()
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Define endpoint
+@app.post("/chat/")
+async def chat(query: Query):
+    question = query.question
 
+    # Step 1: Vector DB search
+    search_results = db.search(question)
+    context = "\n".join([res["content"] for res in search_results])
 
-@app.get("/")
-def read_root():
-    return {"message": "RAG chatbot backend is running with Kalpokoch/QuantizedTinyLlama and ChromaDB!"}
+    # Step 2: Build prompt
+    prompt = f"Context:\n{context}\n\nQuestion: {question}\nAnswer:"
 
+    # Step 3: Tokenize and generate
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
 
-class ChatRequest(BaseModel):
-    question: str
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+    # Optionally strip out the prompt from the output
+    final_answer = answer.split("Answer:")[-1].strip()
 
-@app.post("/chat")
-def chat(request: ChatRequest):
-    question = request.question.strip()
-    if not question:
-        return {"response": "Please ask a question."}
-
-    # --- 3. RAG Retrieval using PolicyVectorDB ---
-    print(f"Searching for context for question: '{question}'")
-    search_results = db.search(query_text=question, top_k=3)
-
-    if not search_results:
-        retrieved_context = "No relevant context found."
-    else:
-        retrieved_context = "\n\n".join([result['text'] for result in search_results])
-
-    print(f"Retrieved Context:\n{retrieved_context[:500]}...")
-
-    # --- 4. Prompt Engineering and Generation ---
-    prompt = (
-        f"<|system|>\nYou are a helpful assistant for NEEPCO policies. "
-        f"Use the following context to answer the user's question. If the context doesn't contain the answer, say that.\n"
-        f"Context:\n{retrieved_context}</s>\n"
-        f"<|user|>\n{question}</s>\n"
-        f"<|assistant|>"
-    )
-
-    try:
-        outputs = pipe(prompt)
-        reply = outputs[0]['generated_text']
-        assistant_reply = reply.split("<|assistant|>")[1].strip()
-        return {"response": assistant_reply}
-    except Exception as e:
-        print(f"Error during model inference: {e}")
-        return {"response": "Sorry, I encountered an error while generating a response."}
+    return {"answer": final_answer}