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
|