File size: 1,463 Bytes
6800266 a13879c 0f79866 6800266 a13879c 6800266 31a1873 3160616 6800266 47de850 6800266 47de850 6800266 47de850 6800266 21d3d7e 6800266 |
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 56 57 58 59 60 61 |
import os
from langchain_huggingface import HuggingFaceEndpoint
from langchain_core.runnables import RunnablePassthrough
from transformers import AutoTokenizer
model_id = "meta-llama/Llama-3.2-3B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
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=model_id
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)
standalone_prompt_formatted =
format_prompt(standalone_prompt)
standalone_chain = standalone_prompt_formatted | llm
generation_chain = rag_prompt_formatted | llm
rag_chain = (
RunnablePassthrough.assign(new_question=standalone_chain)
| {
'context': lambda x:
format_context(search(x['new_question'])),
'standalone_question': lambda x: x['new_question']
}
| generation_chain
).with_types(input_type=schemas.RagInput)
|