# !pip uninstall -y transformers sentence-transformers peft

# !pip install transformers==4.41.2
# !pip install sentence-transformers==2.7.0
# !pip install peft==0.10.0
# !pip install faiss-cpu

from sentence_transformers import SentenceTransformer
import faiss
import numpy as np
from transformers import pipeline


text = """
Section 302: Punishment for murder – Whoever commits murder shall be punished with death or life imprisonment.
Section 378: Theft – Dishonestly taking property without consent.
Section 405: Criminal breach of trust – Misusing entrusted property.
"""


chunks = [text[i:i+200] for i in range(0, len(text), 200)]


model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = model.encode(chunks)

index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(np.array(embeddings))


generator = pipeline("text-generation", model="gpt2")


def chatbot(query):
    query_vec = model.encode([query])
    _, indices = index.search(query_vec, 2)

    context = " ".join([chunks[i] for i in indices[0]])

    prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
    result = generator(prompt, max_length=100)

    return result[0]['generated_text']


while True:
    q = input("You: ")
    if q.lower() == "exit":
        break
    print("Bot:", chatbot(q))