import torch

from transformers import AutoModelForCausalLM, AutoTokenizer

from rag_pipeline import GovernmentProcurementRAGAssistant

MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

assistant = GovernmentProcurementRAGAssistant(
    chunks_path="government_procurement_chunks.csv",
    tokenizer=tokenizer,
    model=model,
    top_k=3,
    max_new_tokens=250
)

response, sources = assistant.answer(
    "What information should a contractor review before bidding on a federal opportunity?"
)

print(response)

print()

print(
    sources[
        [
            "award_id",
            "recipient",
            "agency",
            "similarity_score"
        ]
    ]
)