Update app.py

app.py

@@ -16,25 +16,50 @@ def init_retriever():
 index = init_pinecone()
 retriever = init_retriever()
 
-def card(thubmnail, title, url, context):
-    return st.markdown(f"""
+def card(thumbnail: str, title: str, urls: list, contexts: list, starts: list, ends: list):
+    meta = [(e, s, u, c) for e, s, u, c in zip(ends, starts, urls, contexts)]
+    meta.sort(reverse=False)
+    text_content = []
+    current_start = 0
+    current_end = 0
+    for end, start, url, context in meta:
+        # reformat seconds to timestamp
+        time = start / 60
+        mins = f"0{int(time)}"[-2:]
+        secs = f"0{int(round((time - int(mins))*60, 0))}"[-2:]
+        timestamp = f"{mins}:{secs}"
+        if start < current_end and start > current_start:
+            # this means it is a continuation of the previous sentence
+            text_content[-1][0] = text_content[-1][0].split(context[:10])[0]
+            text_content.append([f"[{timestamp}] {context.capitalize()}", url])
+        else:
+            text_content.append(["xxLINEBREAKxx", ""])
+            text_content.append([f"[{timestamp}] {context}", url])
+        current_start = start
+        current_end = end
+    html_text = ""
+    for text, url in text_content:
+        if text == "xxLINEBREAKxx":
+            html_text += "<br>"
+        else:
+            html_text += f"<small><a href={url}>{text.strip()}... </a></small>"
+    print(html_text)
+    html = f"""
     <div class="container-fluid">
         <div class="row align-items-start">
             <div class="col-md-4 col-sm-4">
-                 <div class="position-relative">
-                     <a href={url}><img src={thubmnail} class="img-fluid" style="width: 192px; height: 106px"></a>
-                 </div>
-             </div>
-             <div  class="col-md-8 col-sm-8">
-                 <a href={url}>{title}</a>
-                 <br>
-                 <span style="color: #808080;">
-                     <small>{context[:200].capitalize()+"...."}</small>
-                 </span>
-             </div>
-        </div>
-     </div>
-        """, unsafe_allow_html=True)
+                <div class="position-relative">
+                    <a href={urls[0]}><img src={thumbnail} class="img-fluid" style="width: 192px; height: 106px"></a>
+                </div>
+            </div>
+            <div  class="col-md-8 col-sm-8">
+                <h2>{title}</h2>
+            </div>
+        <div>
+            {html_text}
+    <br><br>
+    """
+    return st.markdown(html, unsafe_allow_html=True)
 
     
 st.write("""
@@ -52,10 +77,35 @@ if query != "":
     xq = retriever.encode([query]).tolist()
     xc = index.query(xq, top_k=5, include_metadata=True)
     
+    results = {}
+    order = []
     for context in xc['matches']:
+        video_id = context['metadata']['url'].split('/')[-1]
+        if video_id not in results:
+            results[video_id] = {
+                'title': context['metadata']['title'],
+                'urls': [f"{context['metadata']['url']}?t={int(context['metadata']['start'])}"],
+                'contexts': [context['metadata']['text']],
+                'starts': [int(context['metadata']['start'])],
+                'ends': [int(context['metadata']['end'])]
+            }
+            order.append(video_id)
+        else:
+            results[video_id]['urls'].append(
+                f"{context['metadata']['url']}?t={int(context['metadata']['start'])}"
+            )
+            results[video_id]['contexts'].append(
+                context['metadata']['text']
+            )
+            results[video_id]['starts'].append(int(context['metadata']['start']))
+            results[video_id]['ends'].append(int(context['metadata']['end']))
+    # now display cards
+    for video_id in order:
         card(
-            f"https://img.youtube.com/vi/{context['metadata']['url'].split('/')[-1]}/maxresdefault.jpg",
-            context['metadata']['title'],
-            f"{context['metadata']['url']}?t={int(context['metadata']['start'])}",
-            context['metadata']['text']
+            thumbnail=f"https://img.youtube.com/vi/{video_id}/maxresdefault.jpg",
+            title=results[video_id]['title'],
+            urls=results[video_id]['urls'],
+            contexts=results[video_id]['contexts'],
+            starts=results[video_id]['starts'],
+            ends=results[video_id]['ends']
         )