Update main.py

main.py

@@ -6,6 +6,7 @@ from fastapi.staticfiles import StaticFiles
 from fastapi.templating import Jinja2Templates
 from dotenv import load_dotenv
 import os
+from transformers import AutoTokenizer
 # Load environment variables
 load_dotenv()
 from gliner import GLiNER
@@ -19,14 +20,19 @@ app.mount("/static", StaticFiles(directory="static"), name="static")
 # Load models
 cache_dir = os.environ.get("MODEL_CACHE_DIR", "/app/cache")  # Fallback to /app/cache
 os.makedirs(cache_dir, exist_ok=True)
+tokenizer = AutoTokenizer.from_pretrained("microsoft/deberta-v3-base", cache_dir=cache_dir)  # Replace with appropriate tokenizer for GLiNER
 gliner_model = GLiNER.from_pretrained("urchade/gliner_medium-v2.1",cache_dir=cache_dir)
 groq_client = Groq(api_key=GROQ_API_KEY)
 
 init_qdrant_collection()
 
 def extract_entities(text):
+    # Tokenize the input text first
+    inputs = tokenizer(text, return_tensors="pt")  # Assuming PyTorch backend
     labels = ["PRODUCT", "ISSUE", "PROBLEM", "SERVICE"]
-    return gliner_model.predict_entities(text, labels)
+    
+    # Predict entities
+    return gliner_model.predict_entities(inputs['input_ids'], labels)
 
 def validate_answer(user_query, retrieved_answer):
     prompt = f"""