DrishtiSharma commited on
Commit
fab9fe6
·
verified ·
1 Parent(s): 99973bd

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
app.py CHANGED
@@ -343,6 +343,8 @@ class DocumentRAG:
343
  return {"retrieved_docs": docs, "question": query}
344
 
345
 
 
 
346
  def rerank_node(self, state: StoryState):
347
  topic = state.story_topic
348
  query = f"Rerank documents based on how well they explain the topic {topic}"
@@ -352,7 +354,7 @@ class DocumentRAG:
352
  if not texts:
353
  return {"reranked_docs": [], "question": query}
354
 
355
- if self.embedding_choice == "Cohere":
356
  rerank_results = self.cohere_client.rerank(
357
  query=query,
358
  documents=texts,
@@ -368,6 +370,7 @@ class DocumentRAG:
368
 
369
 
370
 
 
371
  def generate_story_node(self, state: StoryState):
372
  context = "\n\n".join(state.reranked_docs)
373
  topic = state.story_topic
 
343
  return {"retrieved_docs": docs, "question": query}
344
 
345
 
346
+
347
+
348
  def rerank_node(self, state: StoryState):
349
  topic = state.story_topic
350
  query = f"Rerank documents based on how well they explain the topic {topic}"
 
354
  if not texts:
355
  return {"reranked_docs": [], "question": query}
356
 
357
+ if self.embedding_choice == "Cohere" and hasattr(self, "cohere_client"):
358
  rerank_results = self.cohere_client.rerank(
359
  query=query,
360
  documents=texts,
 
370
 
371
 
372
 
373
+
374
  def generate_story_node(self, state: StoryState):
375
  context = "\n\n".join(state.reranked_docs)
376
  topic = state.story_topic