Tri4 commited on
Commit
7c5a24d
·
verified ·
1 Parent(s): ca6fb15

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +10 -10
main.py CHANGED
@@ -1,7 +1,7 @@
1
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
2
  from flask import Flask, request, jsonify
3
- import os
4
  import torch
 
5
 
6
  app = Flask(__name__)
7
 
@@ -10,22 +10,23 @@ print("Hello welcome to Sema AI", flush=True) # Flush to ensure immediate outpu
10
  @app.route("/")
11
  def hello():
12
  return "hello 🤗, Welcome to Sema AI Chat Service."
13
-
14
  # Get Hugging Face credentials from environment variables
15
- email = os.getenv('HF_EMAIL')
16
- password = os.getenv('HF_PASS')
17
- GEMMA_TOKEN = os.getenv("GEMMA_TOKEN")
18
 
19
- if not (email and password and GEMMA_TOKEN):
20
- print("No dependencies", flush=True)
21
 
22
  model_id = "google/gemma-2-2b-it"
23
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
24
- tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
25
  model = AutoModelForCausalLM.from_pretrained(
26
  model_id,
27
  device_map="auto",
28
- torch_dtype=torch.float16
 
29
  )
30
 
31
  app_pipeline = pipeline(
@@ -44,7 +45,6 @@ def generate_text():
44
  top_k = data.get("top_k", 50)
45
  top_p = data.get("top_p", 0.95)
46
 
47
- # Generate text using the pipeline
48
  try:
49
  outputs = app_pipeline(
50
  prompt,
 
1
  from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
2
  from flask import Flask, request, jsonify
 
3
  import torch
4
+ import os
5
 
6
  app = Flask(__name__)
7
 
 
10
  @app.route("/")
11
  def hello():
12
  return "hello 🤗, Welcome to Sema AI Chat Service."
13
+
14
  # Get Hugging Face credentials from environment variables
15
+ HF_TOKEN = os.getenv('HF_TOKEN')
 
 
16
 
17
+ if not HF_TOKEN:
18
+ print("Missing Hugging Face token", flush=True)
19
 
20
  model_id = "google/gemma-2-2b-it"
21
  device = "cuda:0" if torch.cuda.is_available() else "cpu"
22
+
23
+ # Load tokenizer and model with authentication token
24
+ tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
25
  model = AutoModelForCausalLM.from_pretrained(
26
  model_id,
27
  device_map="auto",
28
+ torch_dtype=torch.float16,
29
+ use_auth_token=HF_TOKEN
30
  )
31
 
32
  app_pipeline = pipeline(
 
45
  top_k = data.get("top_k", 50)
46
  top_p = data.get("top_p", 0.95)
47
 
 
48
  try:
49
  outputs = app_pipeline(
50
  prompt,