minar09 commited on
Commit
58a106e
·
verified ·
1 Parent(s): 1934471

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +12 -12
main.py CHANGED
@@ -6,8 +6,8 @@ from pathlib import Path
6
  from typing import List, Dict, Optional
7
  from dataclasses import dataclass
8
  from fastapi.encoders import jsonable_encoder
9
- from sentence_transformers import SentenceTransformer
10
- from llama_cpp import Llama
11
 
12
  # Fix: Dynamically adjust the module path if magic_pdf is in a non-standard location
13
  try:
@@ -42,17 +42,17 @@ class PDFProcessor:
42
  self.output_dir.mkdir(exist_ok=True)
43
 
44
  def _initialize_emb_model(self, model_name):
45
- try:
46
- model = SentenceTransformer("sentence-transformers/" + model_name)
47
- model.save('models/'+ model_name)
48
- return model
49
- except:
50
- # Load model directly
51
- from transformers import AutoTokenizer, AutoModel
52
 
53
- tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
54
- model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
55
- return model
56
 
57
  def _initialize_llm(self, model_name):
58
  """Initialize LLM with automatic download if needed"""
 
6
  from typing import List, Dict, Optional
7
  from dataclasses import dataclass
8
  from fastapi.encoders import jsonable_encoder
9
+ # from sentence_transformers import SentenceTransformer
10
+ # from llama_cpp import Llama
11
 
12
  # Fix: Dynamically adjust the module path if magic_pdf is in a non-standard location
13
  try:
 
42
  self.output_dir.mkdir(exist_ok=True)
43
 
44
  def _initialize_emb_model(self, model_name):
45
+ # try:
46
+ # model = SentenceTransformer("sentence-transformers/" + model_name)
47
+ # model.save('models/'+ model_name)
48
+ # return model
49
+ # except:
50
+ # Load model directly
51
+ from transformers import AutoTokenizer, AutoModel
52
 
53
+ tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
54
+ model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
55
+ return model
56
 
57
  def _initialize_llm(self, model_name):
58
  """Initialize LLM with automatic download if needed"""