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"
        ]
    ]
)