# src/summarizer.py

from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from custom_llm import HuggingFaceLLMWrapper


def setup_summarizer():
    llm = HuggingFaceLLMWrapper(model_name="meta-llama/Llama-3.2-3B", device='cuda')  # Change to 'cpu' if necessary
    prompt = PromptTemplate(
        input_variables=["text"],
        template="Summarize the following information concisely:\n\n{text}"
    )
    summarizer = LLMChain(llm=llm, prompt=prompt)

    def summarize_text(text):
        summary = summarizer.run(text)
        return summary

    return summarize_text


if __name__ == "__main__":
    summarizer = setup_summarizer()
    sample_text = "Paracetamol is used to treat pain and fever. It works by inhibiting the production of prostaglandins in the brain..."
    summary = summarizer(sample_text)
    print(f"Summary: {summary}")