new app

galis_app.py

@@ -1,20 +1,33 @@
 from pathlib import Path
 import streamlit as st
-import torch.nn.functional as F
-from predictor.link_predictor import (
-    prepare_system,
-    get_citation_predictions,
-    abstract_to_vector,
-    format_top_k_predictions,
-)
+from dataset.ogbn_link_pred_dataset import OGBNLinkPredDataset
+from model.paper_similarity import PaperSimilarityFinder
 from llm.related_work_generator import generate_related_work
 
-MODEL_PATH = Path("model.pth")
-
 
 @st.cache_resource
-def load_prediction_system(model_path):
-    return prepare_system(model_path)
+def load_similarity_finder():
+    """Ładuje i inicjalizuje PaperSimilarityFinder."""
+    dataset = OGBNLinkPredDataset()
+    model_name = "all-mpnet-base-v2"
+    embeddings_dir = Path("embeddings_cache")
+
+    similarity_finder = PaperSimilarityFinder(
+        dataset,
+        method="sentence_transformer",
+        model_name=model_name,
+        embeddings_cache_path=embeddings_dir,
+    )
+    return similarity_finder, dataset
+
+
+def format_top_k_predictions_from_similarity(similar_papers: list) -> str:
+    """Formatuje listę podobnych artykułów do wyświetlenia w Markdown."""
+    markdown_list = []
+    for i, (idx, score, text) in enumerate(similar_papers):
+        title = text.split('\n')[0].strip()
+        markdown_list.append(f"{i + 1}. **{title}** (Similarity: {score:.4f})")
+    return "\n".join(markdown_list)
 
 
 def app():
@@ -30,7 +43,7 @@ def app():
     if "abstract_text" not in st.session_state:
         st.session_state.abstract_text = ""
 
-    gcn_model, st_model, dataset, z_all = load_prediction_system(MODEL_PATH)
+    similarity_finder, dataset = load_similarity_finder()
 
     col1, col2 = st.columns(2, gap="large")
 
@@ -91,32 +104,26 @@ def app():
                 related_work_placeholder.empty()
 
                 with st.spinner("Analyzing abstract and predicting references..."):
-                    new_vector = abstract_to_vector(
-                        abstract_input, abstract_title, st_model
-                    )
-
-                    probabilities = get_citation_predictions(
-                        vector=F.normalize(new_vector.view(1, -1), p=2, dim=1),
-                        model=gcn_model,
-                        z_all=z_all,
-                        num_nodes=dataset.data.num_nodes,
-                    )
-                    references = format_top_k_predictions(
-                        probabilities, dataset, top_k=num_citations
+                    similar_papers = similarity_finder.find_similar_papers(
+                        title=abstract_title,
+                        abstract=abstract_input,
+                        top_k=num_citations
                     )
+                    references = format_top_k_predictions_from_similarity(similar_papers)
                     st.session_state.references = references
 
                 with references_placeholder.container():
                     st.header("Suggested References")
                     with st.container(height=200):
-                        st.markdown(st.session_state.references)
+                        st.markdown(st.session_state.references, unsafe_allow_html=True)
 
                 with related_work_placeholder.container():
                     with st.spinner("Generating related work section..."):
+                        # Upewnij się, że funkcja generate_related_work akceptuje ten format
                         related_work = generate_related_work(
                             st.session_state.abstract_title,
                             st.session_state.abstract_text,
-                            st.session_state.references
+                            st.session_state.references,
                         )
                         st.session_state.related_work = related_work
 
@@ -124,14 +131,14 @@ def app():
         with references_placeholder.container():
             st.header("Suggested References")
             with st.container(height=200):
-                st.markdown(st.session_state.references)
+                st.markdown(st.session_state.references, unsafe_allow_html=True)
 
     if st.session_state.related_work:
         with related_work_placeholder.container():
             st.header("Suggested Related Works")
             with st.container(height=200):
-                st.markdown(st.session_state.related_work)
+                st.markdown(st.session_state.related_work, unsafe_allow_html=True)
 
 
 if __name__ == "__main__":
-    app()
+    app()