File size: 1,309 Bytes
36118aa
d038098
36118aa
d038098
 
5d0492c
d038098
 
275d03f
 
d038098
 
 
 
 
 
 
 
275d03f
d038098
 
36118aa
d038098
196de9f
 
d038098
 
5d0492c
 
d038098
 
 
 
 
 
 
 
 
 
 
275d03f
d038098
275d03f
 
 
d038098
275d03f
d038098
 
 
 
 
 
 
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
43
44
45
46
47
48
49
50
51
52
53
54
55
import os
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.runnables import RunnablePassthrough
import schemas
from prompts import (
    raw_prompt,
    raw_prompt_formatted, 
    history_prompt_formatted, 
    standalone_prompt_formatted,
    rag_prompt_formatted,
    format_context,
    tokenizer
)
from data_indexing import DataIndexer


data_indexer = DataIndexer()


llm = HuggingFaceEndpoint(
    repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
    huggingfacehub_api_token=os.environ['HF_TOKEN'],
    max_new_tokens=512,
    stop_sequences=[tokenizer.eos_token],
    streaming=True,
)

simple_chain = (raw_prompt | llm).with_types(input_type=schemas.UserQuestion)

formatted_chain = (
    raw_prompt_formatted
    | llm
).with_types(input_type=schemas.UserQuestion)

history_chain = (
    history_prompt_formatted
    | llm
).with_types(input_type=schemas.HistoryInput)

rag_chain = (
    RunnablePassthrough.assign(new_question=standalone_prompt_formatted | llm)
    | {
        'context': lambda x: format_context(data_indexer.search(x['new_question'], hybrid_search=x['hybrid_search'])),
        'standalone_question': lambda x: x['new_question'],
        'test': lambda x : print(x)
    }
    | rag_prompt_formatted
    | llm
).with_types(input_type=schemas.RagInput)