-
Notifications
You must be signed in to change notification settings - Fork 2
/
Mamba.py
99 lines (75 loc) · 3.21 KB
/
Mamba.py
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
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
from langchain_community.embeddings.ollama import OllamaEmbeddings
from langchain_community.vectorstores.chroma import Chroma
import os
from dotenv import load_dotenv
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain_mistralai.chat_models import ChatMistralAI
from langchain_core.prompts import ChatPromptTemplate
# # Create embeddingsclear
embeddings = OllamaEmbeddings(model="nomic-embed-text", show_progress=False)
# embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
db = Chroma(persist_directory="./db-mawared",
embedding_function=embeddings)
# Create retriever
retriever = db.as_retriever(
search_type="similarity",
search_kwargs= {"k": 3}
)
load_dotenv()
os.environ["MISTRAL_API_KEY"] = os.getenv("MISTRAL_API_KEY")
llm = model = ChatMistralAI(model="open-codestral-mamba")
# Create prompt template
template = """
You are an expert assistant specializing in the Mawared HR System. Your role is to answer user questions based strictly on the provided context. If the context is insufficient, ask clarifying questions to gather more information.
Guidelines:
1. Use only the provided context to generate answers.
2. Be concise and direct.
3. If the context is insufficient, ask relevant follow-up questions instead of speculating.
4. Present answers in numbered steps when appropriate.
When responding to a question, follow these steps:
1. Analyze the Question
- Carefully read and comprehend the context and details.
- Break down the question into specific sub-questions if necessary.
- Identify key elements and potential areas requiring clarification.
2. Formulate Response
- Gather relevant information from the provided context.
- Develop a clear and concise answer based on the available information.
- Use analogies, metaphors, or examples to illustrate points when helpful.
3. Verify and Refine
- Ensure your response directly addresses the user's question.
- Check that all information comes from the provided context.
- Identify any gaps in information that require follow-up questions.
4. Present the Answer
- Provide a clear, step-by-step response when appropriate.
- Use an engaging and accessible tone.
- Acknowledge any limitations in the available information.
- Ask follow-up questions if the context is insufficient.
Context:
{context}
Question: {question}
Answer:
"""
prompt = ChatPromptTemplate.from_template(template)
# Create the RAG chain using LCEL with prompt printing and streaming output
rag_chain = (
{"context": retriever, "question": RunnablePassthrough()}
| prompt
| llm
| StrOutputParser()
)
# Function to ask questions
def ask_question(question):
print("Answer:\t", end=" ", flush=True)
for chunk in rag_chain.stream(question):
print(chunk, end="", flush=True)
print("\n")
# Example usage
if __name__ == "__main__":
while True:
user_question = input("Ask a question (or type 'quit' to exit): ")
if user_question.lower() == 'quit':
break
answer = ask_question(user_question)
print("\nFull answer received.\n")