main.py

from sentence_transformers import SentenceTransformer
from llama_cpp import Llama
import faiss
import os


embedder = SentenceTransformer('all-MiniLM-L6-v2')


def load_text(file_path):
    with open(file_path, 'r', encoding='utf-8') as f:
        return f.read()


def chunk_text(text, chunk_size=500):
    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]


def create_faiss_index(chunks):
    embeddings = embedder.encode(chunks)
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index, chunks


def search_index(query, index, chunks, top_k=3):
    q_vec = embedder.encode([query])
    distances, indices = index.search(q_vec, top_k)
    return [chunks[i] for i in indices[0]]


llm = Llama(
    model_path="models/llama-2-7b.Q4_K_M.gguf",
    n_ctx=2048,
    n_threads=6,
    verbose=False
)


def generate_answer(context, question):
    prompt = f"""You are a helpful assistant. Use the following context to answer the question.

Context:
{context}

Question: {question}
Answer:"""
    output = llm(prompt, max_tokens=256, stop=["\n\n"])
    return output["choices"][0]["text"].strip()

def main():
    file_path = 'docs/nsx.txt'
    if not os.path.exists(file_path):
        print(f"File not found: {file_path}")
        return

    print("Reading file and building index...")
    text = load_text(file_path)
    chunks = chunk_text(text)
    index, stored_chunks = create_faiss_index(chunks)

    print(f"Ready! {len(chunks)} chunks indexed.")

    while True:
        query = input("\nAsk your question (or 'exit'): ")
        if query.lower() == 'exit':
            break

        top_chunks = search_index(query, index, stored_chunks)
        context = "\n".join(top_chunks)
        answer = generate_answer(context, query)

        print(f"\n Answer:\n{answer}\n")

if __name__ == '__main__':
    main()