File size: 1,185 Bytes
36118aa
d038098
36118aa
d038098
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
36118aa
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
import os
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.runnables import RunnablePassthrough
import schemas
from prompts import (
    raw_prompt_formatted, 
    history_prompt_formatted, 
    question_prompt_formatted,
    context_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]
)

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 = (
    {
      'question': question_prompt_formatted | llm,
      'hybrid_search': RunnablePassthrough()   
    }
    | {
        'context': lambda x: format_context(data_indexer.search(x['question'], hybrid_search=x['hybrid_search'])),
        'standalone_question': lambda x: x['question']
    }
    | context_prompt_formatted
    | llm
).with_types(input_type=schemas.RagInput)