main.py

"""
main.py
Optional: Text-only RAG demo using SentenceTransformers, FAISS, and vllm.
"""

from sentence_transformers import SentenceTransformer
from vllm import LLM, SamplingParams
import faiss
import numpy as np

# Load embedding model and LLM
embedder = SentenceTransformer('all-MiniLM-L6-v2')
llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")
params = SamplingParams(temperature=0.7)

# Sample docs
docs = [
    "Paris is the capital of France.",
    "Joe Biden is the president of the US.",
    "AI is rapidly evolving."
]

# Embed and store in FAISS
doc_embeds = embedder.encode(docs, normalize_embeddings=True)
index = faiss.IndexFlatIP(doc_embeds.shape[1])
index.add(np.array(doc_embeds))

# Ask a question
query = "Who is the president of france?"
query_embed = embedder.encode([query], normalize_embeddings=True)
D, I = index.search(np.array(query_embed), k=1)

# Retrieve top doc
context = docs[I[0][0]]

# RAG-style prompt
prompt = f"Context: {context}\n\nQuestion: {query}\nAnswer:"

# Generate response
output = llm.generate([prompt], params)[0]
print("RAG Answer:", output.outputs[0].text)