File size: 841 Bytes
36787cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 | 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"
]
]
) |