Update app.py
Browse files
app.py
CHANGED
@@ -31,6 +31,13 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
|
|
31 |
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
|
32 |
# Database configuration
|
33 |
DATABASE_PATH = '/tmp/chat_database.db'
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
34 |
|
35 |
# Initialize LangChain with Ollama LLM
|
36 |
if hf_token:
|
@@ -40,12 +47,13 @@ else:
|
|
40 |
model_name = "microsoft/phi-4"
|
41 |
|
42 |
try:
|
43 |
-
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
44 |
model = AutoModelForCausalLM.from_pretrained(
|
45 |
model_name,
|
46 |
torch_dtype=torch.float16,
|
47 |
device_map="auto",
|
48 |
-
load_in_8bit=True
|
|
|
49 |
)
|
50 |
|
51 |
# Create pipeline
|
|
|
31 |
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
|
32 |
# Database configuration
|
33 |
DATABASE_PATH = '/tmp/chat_database.db'
|
34 |
+
CACHE_DIR = "/tmp/huggingface_cache"
|
35 |
+
MODEL_CACHE_DIR = "/tmp/model_cache"
|
36 |
+
os.makedirs(CACHE_DIR, exist_ok=True)
|
37 |
+
os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
|
38 |
+
os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
|
39 |
+
os.environ['HF_HOME'] = CACHE_DIR
|
40 |
+
os.environ['HF_DATASETS_CACHE'] = CACHE_DIR
|
41 |
|
42 |
# Initialize LangChain with Ollama LLM
|
43 |
if hf_token:
|
|
|
47 |
model_name = "microsoft/phi-4"
|
48 |
|
49 |
try:
|
50 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=CACHE_DIR)
|
51 |
model = AutoModelForCausalLM.from_pretrained(
|
52 |
model_name,
|
53 |
torch_dtype=torch.float16,
|
54 |
device_map="auto",
|
55 |
+
load_in_8bit=True,
|
56 |
+
cache_dir=MODEL_CACHE_DIR
|
57 |
)
|
58 |
|
59 |
# Create pipeline
|