Spaces:
Sleeping
Sleeping
Update main.py
Browse files
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 |
-
|
16 |
-
password = os.getenv('HF_PASS')
|
17 |
-
GEMMA_TOKEN = os.getenv("GEMMA_TOKEN")
|
18 |
|
19 |
-
if not
|
20 |
-
print("
|
21 |
|
22 |
model_id = "google/gemma-2-2b-it"
|
23 |
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
24 |
-
|
|
|
|
|
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,
|