updates app

app.py

@@ -1,12 +1,11 @@
 import asyncio
 import time
 from fastapi import FastAPI, Request, HTTPException
-# # from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse
 from starlette.status import HTTP_504_GATEWAY_TIMEOUT
 from neural_searcher import NeuralSearcher
-# from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import login
-from config import HUGGING_FACE_API_KEY,COLLECTION_NAME, ALLOWED_ORIGINS, API_KEY
+from config import HUGGING_FACE_API_KEY,COLLECTION_NAME, API_KEY
 
 login(HUGGING_FACE_API_KEY)
 
@@ -14,53 +13,33 @@ app = FastAPI()
 
 neural_searcher = NeuralSearcher(collection_name=COLLECTION_NAME)
 
-# REQUEST_TIMEOUT_ERROR = 30
+REQUEST_TIMEOUT_ERROR = 30
 
-# ALLOWED_ORIGINS = [ALLOWED_ORIGINS]
-# ALLOWED_API_KEY = API_KEY 
+ALLOWED_API_KEY = API_KEY 
 
 
 @app.get("/api/search")
 async def search(q: str):
     data =  await neural_searcher.search(text=q)
     return data
-    
 
-# app.add_middleware(
-#     CORSMiddleware,
-#     allow_origins=ALLOWED_ORIGINS,
-#     allow_credentials=True,
-#     allow_methods=["GET"],
-#     allow_headers=["*"],
-# )
+@app.middleware("http")
+async def security_middleware(request: Request, call_next):
+    api_key = request.headers.get("X-API-KEY", "")
 
-# @app.middleware("http")
-# async def security_middleware(request: Request, call_next):
-#     referer = request.headers.get("referer", "")
-#     origin = request.headers.get("origin", "")
-#     user_agent = request.headers.get("user-agent", "")
-#     api_key = request.headers.get("X-API-KEY", "")
+    if api_key != ALLOWED_API_KEY:
+        raise HTTPException(status_code=403, detail="Access denied.")
 
-    
-#     if not (referer.startswith(ALLOWED_ORIGINS[0]) or origin.startswith(ALLOWED_ORIGINS[0])):
-#         raise HTTPException(status_code=403, detail="Access denied: Invalid source")
+    return await call_next(request)
 
-#     if not user_agent or "Mozilla" not in user_agent:
-#         raise HTTPException(status_code=403, detail="Access denied: Suspicious client")
+@app.middleware("http")
+async def timeout_middleware(request: Request, call_next):
+    try:
+        start_time = time.time()
+        return await asyncio.wait_for(call_next(request), timeout=REQUEST_TIMEOUT_ERROR)
 
-#     if api_key != ALLOWED_API_KEY:
-#         raise HTTPException(status_code=403, detail="Access denied: Invalid API Key")
-
-#     return await call_next(request)
-
-# @app.middleware("http")
-# async def timeout_middleware(request: Request, call_next):
-#     try:
-#         start_time = time.time()
-#         return await asyncio.wait_for(call_next(request), timeout=REQUEST_TIMEOUT_ERROR)
-
-#     except asyncio.TimeoutError:
-#         process_time = time.time() - start_time
-#         return JSONResponse({'detail': 'Request processing time excedeed limit',
-#                              'processing_time': process_time},
-#                             status_code=HTTP_504_GATEWAY_TIMEOUT)
+    except asyncio.TimeoutError:
+        process_time = time.time() - start_time
+        return JSONResponse({'detail': 'Request processing time excedeed limit',
+                             'processing_time': process_time},
+                            status_code=HTTP_504_GATEWAY_TIMEOUT)

ner.py

@@ -1,96 +0,0 @@
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-from config import NER_MODEL
-import torch
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer = AutoTokenizer.from_pretrained(NER_MODEL, use_auth_token=True)
-model = AutoModelForTokenClassification.from_pretrained(NER_MODEL, use_auth_token=True).to(device)
-
-id_to_label = {
-    0: 'O',
-    1: 'B-COURT',
-    2: 'B-DATE',
-    3: 'B-DECISION',
-    4: 'B-LAW',
-    5: 'B-MONEY',
-    6: 'B-OFFICIAL GAZZETE',
-    7: 'B-PERSON',
-    8: 'B-REFERENCE',
-    9: 'I-COURT',
-    10: 'I-LAW',
-    11: 'I-MONEY',
-    12: 'I-OFFICIAL GAZZETE',
-    13: 'I-PERSON',
-    14: 'I-REFERENCE'
-}
-
-def perform_ner(text):
-    try:
-        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
-        with torch.no_grad():
-            outputs = model(**inputs)
-        logits = outputs.logits
-        predictions = torch.argmax(logits, dim=2).squeeze().tolist()
-
-    except RuntimeError as e:
-        if "CUDA out of memory" in str(e):
-            print("Switching to CPU due to memory constraints.")
-            inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-            with torch.no_grad():
-                outputs = model.cpu()(**inputs)  # Run model on CPU
-            logits = outputs.logits
-            predictions = torch.argmax(logits, dim=2).squeeze().tolist()
-        else:
-            raise e 
-    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"].squeeze())
-    labels = [id_to_label[pred] for pred in predictions]
-
-    results = [
-        (token, label)
-        for token, label in zip(tokens, labels)
-        if token not in tokenizer.all_special_tokens
-    ]
-    return results
-
-text = ""
-
-def merge_entities(token_label_pairs):
-    merged_words, merged_labels = [], []
-    current_word, current_label = "", None
-
-    for token, label in token_label_pairs:
-        if token.startswith("##"):
-            current_word += token[2:]
-        else:
-            if current_word:
-                merged_words.append(current_word)
-                merged_labels.append(current_label)
-
-            current_word, current_label = token, label
-
-    if current_word:
-        merged_words.append(current_word)
-        merged_labels.append(current_label)
-
-    final_words, final_labels = [], []
-    
-    for i, (word, label) in enumerate(zip(merged_words, merged_labels)):
-        if final_labels and (
-            label == final_labels[-1] or
-            (label.startswith("I-") and final_labels[-1].endswith(label[2:])) or
-            (label.startswith("B-") and final_labels[-1].endswith(label[2:]))  
-        ):  
-            
-            final_words[-1] += " " + word  
-        else:
-            final_words.append(word)
-            final_labels.append(label)
-
-    return final_words, final_labels
-
-results = perform_ner(text)
-
-words,labels = merge_entities(results)
-
-for i,b in zip(words,labels):
-    print(i + " ### " + b)

senatus_client.py

@@ -1,136 +0,0 @@
-import json
-import uuid
-import numpy as np
-import os
-from huggingface_hub import login
-from fastembed import SparseTextEmbedding,LateInteractionTextEmbedding
-from qdrant_client import QdrantClient, models
-from sentence_transformers import SentenceTransformer
-from tqdm import tqdm
-from huggingface_hub import login
-from config import HUGGING_FACE_API_KEY, DENSE_MODEL, SPARSE_MODEL, LATE_INTERACTION_MODEL, QDRANT_URL, QDRANT_API_KEY, COLLECTION_NAME
-
-login(HUGGING_FACE_API_KEY)
-
-folder_path = 'data'
-
-dense_model = SentenceTransformer(DENSE_MODEL)
-sparse_model = SparseTextEmbedding(SPARSE_MODEL)
-# late_interaction_embedding_model = LateInteractionTextEmbedding(LATE_INTERACTION_MODEL)
-
-data = []
-for filename in os.listdir(folder_path):
-    if filename.endswith('.json'):
-        file_path = os.path.join(folder_path, filename)
-        with open(file_path,encoding='utf-8') as f:
-            data = json.load(f)
-
-
-client = QdrantClient(QDRANT_URL,api_key=QDRANT_API_KEY)
-
-
-data_array = np.array(data)
-
-split_data = np.array_split(data_array, 1000)
-
-collection_name = COLLECTION_NAME
-
-for local_data in split_data:
-
-  payload = []
-  documents = []
-  for obj in local_data:
-      documents.append(obj["tekst"])
-      payload.append(obj)
-
-  sparse_embeddings = list(
-      tqdm(
-          sparse_model.passage_embed(doc for doc in documents),
-          total=len(documents),
-          desc="🔨 Encoding Sparse Embeddings"
-      )
-  )
-
-#   late_interaction_embeddings = list(
-#       tqdm(
-#           late_interaction_embedding_model.passage_embed(doc for doc in documents),
-#           total=len(documents),
-#           desc="🔨 Encoding Late Interaction Embeddings"
-#       )
-#   )
-
-  dense_embeddings = dense_model.encode(documents, show_progress_bar=True, device="cuda")
-
-  existing_collections = client.get_collections().collections
-  collection_names = [col.name for col in existing_collections]
-
-  if collection_name not in collection_names:
-      client.create_collection(
-          collection_name=collection_name,
-          vectors_config={
-              DENSE_MODEL: models.VectorParams(
-                  size=len(dense_embeddings[0]),
-                  distance=models.Distance.COSINE,
-                  on_disk=True
-              ),
-            #   LATE_INTERACTION_MODEL: models.VectorParams(
-            #       size=len(late_interaction_embeddings[0][0]),
-            #       distance=models.Distance.COSINE,
-            #       multivector_config=models.MultiVectorConfig(
-            #           comparator=models.MultiVectorComparator.MAX_SIM,
-            #       ),
-            #       hnsw_config=models.HnswConfigDiff(
-            #         m=0,  # Disable HNSW graph creation
-            #       ),
-            #       on_disk=True
-            #   ),
-          },
-          sparse_vectors_config={
-              SPARSE_MODEL: models.SparseVectorParams(
-                  modifier=models.Modifier.IDF,
-              ),
-          },
-          quantization_config=models.ScalarQuantization(
-              scalar=models.ScalarQuantizationConfig(
-              type=models.ScalarType.INT8,
-              always_ram=True
-              )
-          ),
-          optimizers_config=models.OptimizersConfigDiff(
-            indexing_threshold=10000,
-          ),
-          shard_number = 4,
-          hnsw_config=models.HnswConfigDiff(on_disk=True),
-      )
-
-  print("🚀 Uploading to qdrant collection: " + collection_name)
-  client.upload_points(
-      collection_name=collection_name,
-      batch_size = 32,
-      parallel = 16,
-      points=[
-          models.PointStruct(
-              id=uuid.uuid4().hex,
-              vector={
-                  DENSE_MODEL: dense_embedding,
-                  SPARSE_MODEL: sparse_embedding.as_object(),
-                #   "answerdotai/answerai-colbert-small-v1":late_interaction_embedding
-              },
-              payload=doc,
-          )
-          for doc, dense_embedding, sparse_embedding in zip(
-              payload, dense_embeddings, sparse_embeddings
-          )
-      ],
-  )
-
-client.create_payload_index(
-    collection_name=collection_name,
-    field_name="dbid",
-    field_schema=models.PayloadSchemaType.INTEGER
-)
-
-client.update_collection(
-    collection_name=collection_name,
-    optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000),
-)