Update app.py

app.py

@@ -41,7 +41,7 @@ recursive_framework_document = metadata_generator(ai_framework_document, "AI Fra
 recursive_blueprint_document = metadata_generator(ai_blueprint_document, "AI Blueprint")
 combined_documents = recursive_framework_document + recursive_blueprint_document
 
-
+"""
 #from transformers import AutoTokenizer, AutoModel
 #import torch
 #embedding = AutoModel.from_pretrained("Cheselle/finetuned-arctic-sentence")
@@ -57,6 +57,43 @@ vectorstore = Qdrant.from_documents(
     location=":memory:",
     collection_name="ai_policy"
 )
+"""
+
+from transformers import AutoTokenizer, AutoModel
+import torch
+from qdrant_client import QdrantClient
+
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("Cheselle/finetuned-arctic-sentence")
+model = AutoModel.from_pretrained("Cheselle/finetuned-arctic-sentence")
+
+# Define a wrapper function for embedding documents
+def embed(documents):
+    inputs = tokenizer(documents, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        return outputs.last_hidden_state.mean(dim=1).numpy()  # Return embeddings
+
+# Initialize Qdrant client (in-memory for testing)
+qdrant_client = QdrantClient(":memory:")
+
+# Create the Qdrant collection
+qdrant_client.recreate_collection(
+    collection_name="ai_policy",
+    vectors_config={"size": 768, "distance": "Cosine"}  # Adjust size based on embedding dimensions
+)
+
+# Example documents
+combined_documents = ["This is document one.", "This is document two."]
+
+# Create vectorstore (with embedding function)
+vectorstore = qdrant_client.upsert(
+    collection_name="ai_policy",
+    points=[
+        {"id": i, "vector": embed([doc])[0], "payload": {"document": doc}}
+        for i, doc in enumerate(combined_documents)
+    ]
+)
 retriever = vectorstore.as_retriever()
 
 ## Generation LLM