Update app.py

app.py

@@ -127,78 +127,74 @@ vectorstore = Chroma.from_documents(
 )
 
 
+import cohere
+from langchain_core.prompts import PromptTemplate
+
+
 class RAGPipeline:
-    def __init__(self, vectorstore, model_name="CohereForAI/aya-expanse-8b", k=6):
+    def __init__(self, vectorstore, api_key, model_name="c4ai-aya-expanse-8b", k=3):
         self.vectorstore = vectorstore
         self.model_name = model_name
         self.k = k
+        self.api_key = api_key
+        self.client = cohere.Client(api_key)  # Initialize the Cohere client
         self.retriever = self.vectorstore.as_retriever(
-            search_type="mmr", search_kwargs={"k": self.k}
+            search_type="mmr", search_kwargs={"k": 3}
         )
         self.prompt_template = PromptTemplate.from_template(self._get_template())
 
-        # Load model and tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, token=token)
-        self.model = AutoModelForCausalLM.from_pretrained(
-            self.model_name, torch_dtype=torch.bfloat16, device_map="auto", token=token
-        )
-
     def _get_template(self):
-        return """\
-        <s>[INST] <<SYS>>
+        return """<s>[INST] <<SYS>>
         أنت مساعد مفيد يقدم إجابات باللغة العربية بناءً على السياق المقدم.
         - أجب فقط باللغة العربية
         - إذا لم تجد إجابة في السياق، قل أنك لا تعرف
         - كن دقيقاً وواضحاً في إجاباتك
+        -جاوب من السياق حصريا
         <</SYS>>
 
         السياق: {context}
 
         السؤال: {question}
         الإجابة: [/INST]\
-        """
+
+"""
 
     def generate_response(self, question):
         retrieved_docs = self._retrieve_documents(question)
         prompt = self._create_prompt(retrieved_docs, question)
-        response = self._generate_response(prompt)
+        response = self._generate_response_cohere(prompt)
         return response
 
     def _retrieve_documents(self, question):
-        start = time.time()
         retrieved_docs = self.retriever.invoke(question)
-        result = {f"doc_{i}": doc.page_content for i, doc in enumerate(retrieved_docs)}
-        end = time.time()
-        time_lapsed = end - start
-        print(f"Time lapsed in Retreival: {time_lapsed}")
-        return result
+        # print("\n=== المستندات المسترجعة ===")
+        # for i, doc in enumerate(retrieved_docs):
+        #     print(f"المستند {i+1}: {doc.page_content}")
+        #     print("==========================\n")
+
+        # دمج النصوص المسترجعة في سياق واحد
+        return " ".join([doc.page_content for doc in retrieved_docs])
 
     def _create_prompt(self, docs, question):
         return self.prompt_template.format(context=docs, question=question)
 
-    def _generate_response(self, prompt):
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
-
-        start = time.time()
-        outputs = self.model.generate(
-            inputs.input_ids,
-            max_new_tokens=1024,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=self.tokenizer.eos_token_id,
+    def _generate_response_cohere(self, prompt):
+        # Call Cohere's generate API
+        response = self.client.generate(
+            model=self.model_name,
+            prompt=prompt,
+            max_tokens=2000,  # Adjust token limit based on requirements
+            temperature=0.3,  # Control creativity
+            stop_sequences=None,
         )
-        end = time.time()
-        time_lapsed = end - start
-        print(f"Time lapsed in Generation: {time_lapsed}")
-
-        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Extract only the assistant's response after [/INST]
-        return response.split("[/INST]")[-1].strip()
-
 
-rag_pipeline = RAGPipeline(vectorstore)
+        if response.generations:
+            return response.generations[0].text.strip()
+        else:
+            raise Exception("No response generated by Cohere API.")
 
+api_key = os.getenv("API_KEY")
+rag_pipeline = RAGPipeline(vectorstore=vectorstore, api_key=api_key)
 
 question = st.text_area("أدخل سؤالك هنا")
 if st.button("Generate Answer"):