roshnn24 commited on
Commit
4766a87
·
verified ·
1 Parent(s): 1e4f197

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +32 -5
app.py CHANGED
@@ -12,21 +12,47 @@ from typing import Dict, List
12
  import sqlite3
13
  from contextlib import contextmanager
14
  import re
15
- import os
16
  from werkzeug.utils import secure_filename
17
 
18
  app = Flask(__name__)
19
 
20
- UPLOAD_FOLDER = 'uploads'
 
 
21
  ALLOWED_EXTENSIONS = {'py'}
22
  app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
23
  os.makedirs(UPLOAD_FOLDER, exist_ok=True)
24
  # Database configuration
25
- DATABASE_PATH = 'chat_database.db'
26
 
27
  # Initialize LangChain with Ollama LLM
28
- llm = Ollama(model="mistral:7b")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
 
30
 
31
  @contextmanager
32
  def get_db_connection():
@@ -606,4 +632,5 @@ def home():
606
 
607
 
608
  if __name__ == "__main__":
609
- app.run(debug=True)
 
 
12
  import sqlite3
13
  from contextlib import contextmanager
14
  import re
 
15
  from werkzeug.utils import secure_filename
16
 
17
  app = Flask(__name__)
18
 
19
+ PORT = int(os.environ.get("PORT", 7860))
20
+
21
+ UPLOAD_FOLDER = '/tmp/uploads' # Change to tmp directory for Spaces
22
  ALLOWED_EXTENSIONS = {'py'}
23
  app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
24
  os.makedirs(UPLOAD_FOLDER, exist_ok=True)
25
  # Database configuration
26
+ DATABASE_PATH = '/tmp/chat_database.db'
27
 
28
  # Initialize LangChain with Ollama LLM
29
+ from transformers import AutoTokenizer, AutoModelForCausalLM
30
+ import torch
31
+
32
+ # Load model and tokenizer
33
+ model_name = "mistralai/Mistral-7B-Instruct-v0.1"
34
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
35
+ model = AutoModelForCausalLM.from_pretrained(model_name)
36
+
37
+ class HuggingFaceLLM:
38
+ def __init__(self, model, tokenizer):
39
+ self.model = model
40
+ self.tokenizer = tokenizer
41
+
42
+ def predict(self, prompt):
43
+ inputs = self.tokenizer(prompt, return_tensors="pt", max_length=2048, truncation=True)
44
+ with torch.no_grad():
45
+ outputs = self.model.generate(
46
+ inputs["input_ids"],
47
+ max_length=2048,
48
+ num_return_sequences=1,
49
+ temperature=0.7,
50
+ pad_token_id=self.tokenizer.eos_token_id
51
+ )
52
+ response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
53
+ return response
54
 
55
+ llm = HuggingFaceLLM(model, tokenizer)
56
 
57
  @contextmanager
58
  def get_db_connection():
 
632
 
633
 
634
  if __name__ == "__main__":
635
+ init_db()
636
+ app.run(host="0.0.0.0", port=PORT)