Callmebowoo-22 commited on
Commit
42c78bd
·
verified ·
1 Parent(s): 3ce3453

Create utils/model.py

Browse files
Files changed (1) hide show
  1. utils/model.py +28 -0
utils/model.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline, AutoModelForTimeSeriesPrediction
2
+ import torch
3
+
4
+ device = "cuda" if torch.cuda.is_available() else "cpu"
5
+
6
+ def predict_umkm(data):
7
+ # Load model TTM
8
+ ttm = AutoModelForTimeSeriesPrediction.from_pretrained(
9
+ "ibm/granite-ttm-r2"
10
+ ).to(device)
11
+
12
+ # Prediksi demand
13
+ inputs = {"values": data['demand'].tolist()}
14
+ demand_pred = ttm.generate(**inputs, max_length=7)
15
+
16
+ # Rekomendasi dengan Chronos-T5
17
+ chronos = pipeline(
18
+ "text-generation",
19
+ model="amazon/chronos-t5-small",
20
+ device=device
21
+ )
22
+ prompt = f"""
23
+ Data UMKM:
24
+ - Prediksi demand: {demand_pred}
25
+ - Stok saat ini: {data['supply'].iloc[-1]}
26
+ Beri rekomendasi dalam 1 kalimat:
27
+ """
28
+ return chronos(prompt, max_length=50)[0]['generated_text']