deepapaikar
commited on
Commit
•
ccc774e
1
Parent(s):
edbe2ee
Update app.py
Browse files
app.py
CHANGED
@@ -9,7 +9,18 @@ torch.set_default_device("cuda")
|
|
9 |
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
10 |
|
11 |
|
12 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
13 |
|
14 |
# pipeline = transformers.pipeline(
|
15 |
# "text-generation",
|
@@ -17,7 +28,7 @@ model = "deepapaikar/katzbot-phi2"
|
|
17 |
# torch_dtype=torch.float16,
|
18 |
# )
|
19 |
|
20 |
-
tokenizer = AutoTokenizer.from_pretrained(
|
21 |
|
22 |
|
23 |
# def predict_answer(question, token=25):
|
@@ -45,7 +56,7 @@ def predict_answer(question, token=25):
|
|
45 |
inputs = {k: v.to(device) for k, v in inputs.items()}
|
46 |
|
47 |
# Use the model directly for inference
|
48 |
-
model.eval()
|
49 |
model.to(device) # Ensure the model is on the correct device
|
50 |
|
51 |
# Generate outputs
|
|
|
9 |
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
10 |
|
11 |
|
12 |
+
# model_name = "deepapaikar/katzbot-phi2"
|
13 |
+
|
14 |
+
# model = AutoModelForCausalLM.from_from_pretrained(model_name)
|
15 |
+
|
16 |
+
|
17 |
+
# Initialize the model and tokenizer
|
18 |
+
model = AutoModelForCausalLM.from_pretrained("deepapaikar/katzbot-phi2",
|
19 |
+
torch_dtype=torch.float16,
|
20 |
+
device_map="auto",
|
21 |
+
trust_remote_code=True)
|
22 |
+
tokenizer = AutoTokenizer.from_pretrained("deepapaikar/katzbot-phi2", trust_remote_code=True)
|
23 |
+
|
24 |
|
25 |
# pipeline = transformers.pipeline(
|
26 |
# "text-generation",
|
|
|
28 |
# torch_dtype=torch.float16,
|
29 |
# )
|
30 |
|
31 |
+
# tokenizer = AutoTokenizer.from_pretrained(model_name)
|
32 |
|
33 |
|
34 |
# def predict_answer(question, token=25):
|
|
|
56 |
inputs = {k: v.to(device) for k, v in inputs.items()}
|
57 |
|
58 |
# Use the model directly for inference
|
59 |
+
model.eval() # Ensure the model is in evaluation mode
|
60 |
model.to(device) # Ensure the model is on the correct device
|
61 |
|
62 |
# Generate outputs
|