minar09 commited on
Commit
79cdd30
·
verified ·
1 Parent(s): 4bc3210

Upload main.py

Browse files
Files changed (1) hide show
  1. main.py +9 -3
main.py CHANGED
@@ -42,8 +42,14 @@ class PDFProcessor:
42
  self.output_dir.mkdir(exist_ok=True)
43
 
44
  def _initialize_emb_model(self, model_name):
45
- model = SentenceTransformer("sentence-transformers/" + model_name)
46
- model.save('models/'+ model_name)
 
 
 
 
 
 
47
  return model
48
 
49
  def _initialize_llm(self, model_name):
@@ -187,4 +193,4 @@ def process_pdf_catalog(pdf_path: str):
187
  return result, "Processing completed successfully!"
188
  except Exception as e:
189
  logger.error(f"Processing failed: {e}")
190
- return {}, "Error processing PDF"
 
42
  self.output_dir.mkdir(exist_ok=True)
43
 
44
  def _initialize_emb_model(self, model_name):
45
+ # model = SentenceTransformer("sentence-transformers/" + model_name)
46
+ # model = SentenceTransformer(model_name)
47
+ # model.save('models/'+ model_name)
48
+ # Load model directly
49
+ from transformers import AutoTokenizer, AutoModel
50
+
51
+ tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
52
+ model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
53
  return model
54
 
55
  def _initialize_llm(self, model_name):
 
193
  return result, "Processing completed successfully!"
194
  except Exception as e:
195
  logger.error(f"Processing failed: {e}")
196
+ return {}, "Error processing PDF"