Update main.py

main.py

@@ -5,6 +5,8 @@ from transformers import AutoTokenizer, AutoModel
 from sklearn.metrics.pairwise import cosine_similarity
 from flask import Flask, request, jsonify
 import torch
+import os
+
 app = Flask(__name__)
 
 n_gpu_layers = 0