src/alternatives_generator.py

# src/alternatives_generator.py

from langchain.chains import RetrievalQA
from custom_llm import HuggingFaceLLMWrapper
from rag_application import load_embeddings_db


def setup_alternatives_generator():
    db = load_embeddings_db()
    llm = HuggingFaceLLMWrapper(model_name="meta-llama/Llama-3.2-3B", device='cuda')  # Change to 'cpu' if necessary
    qa = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=db.as_retriever()
    )

    def suggest_alternatives(medication):
        query = f"Suggest alternative medications to {medication}."
        response = qa.run(query)
        return response

    return suggest_alternatives


if __name__ == "__main__":
    generator = setup_alternatives_generator()
    medication = "Paracetamol"
    alternatives = generator(medication)
    print(f"Alternatives to {medication}: {alternatives}")