Spaces:
Sleeping
Sleeping
to use lcel in refusal_chain.py
Browse files- refusal_chain.py +6 -8
refusal_chain.py
CHANGED
@@ -1,20 +1,18 @@
|
|
1 |
# refusal_chain.py
|
2 |
import os
|
3 |
-
from langchain.chains import LLMChain
|
4 |
from langchain_groq import ChatGroq
|
5 |
-
|
6 |
from prompts import refusal_prompt
|
7 |
|
8 |
def get_refusal_chain() -> LLMChain:
|
9 |
"""
|
10 |
Builds the refusal chain for out-of-scope queries.
|
11 |
"""
|
12 |
-
|
13 |
model="Gemma2-9b-It",
|
14 |
-
groq_api_key=os.environ["GROQ_API_KEY"]
|
15 |
-
)
|
16 |
-
chain = LLMChain(
|
17 |
-
llm=chat_groq_model,
|
18 |
-
prompt=refusal_prompt
|
19 |
)
|
|
|
|
|
20 |
return chain
|
|
|
1 |
# refusal_chain.py
|
2 |
import os
|
3 |
+
# from langchain.chains import LLMChain
|
4 |
from langchain_groq import ChatGroq
|
5 |
+
from langchain_core.output_parsers import StrOutputParser
|
6 |
from prompts import refusal_prompt
|
7 |
|
8 |
def get_refusal_chain() -> LLMChain:
|
9 |
"""
|
10 |
Builds the refusal chain for out-of-scope queries.
|
11 |
"""
|
12 |
+
chat_groq_model = ChatGroq(
|
13 |
model="Gemma2-9b-It",
|
14 |
+
groq_api_key=os.environ["GROQ_API_KEY"] # must be set in environment
|
|
|
|
|
|
|
|
|
15 |
)
|
16 |
+
output_parser=StrOutputParser()
|
17 |
+
chain=refusal_prompt|chat_groq_model|output_parser
|
18 |
return chain
|