Update utils/model.py

utils/model.py

@@ -1,28 +1,46 @@
 from transformers import pipeline, AutoModelForTimeSeriesPrediction
 import torch
+import numpy as np
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def predict_umkm(data):
-    # Load model TTM
-    ttm = AutoModelForTimeSeriesPrediction.from_pretrained(
-        "ibm/granite-ttm-r2"
-    ).to(device)
+    try:
+        # ========== [1] GRANITE-TTM FORECASTING ==========
+        ttm_model = AutoModelForTimeSeriesPrediction.from_pretrained(
+            "ibm/granite-ttm-r2"
+        ).to(device)
+        
+        # Format input khusus untuk TTM
+        inputs = {
+            "values": torch.tensor(data['demand'].tolist()).float().unsqueeze(0).to(device),
+            "attention_mask": torch.ones(len(data['demand'])).unsqueeze(0).to(device)
+        }
+        
+        # Generate prediksi
+        with torch.no_grad():
+            demand_pred = ttm_model.generate(**inputs, max_length=7)
+        
+        # Konversi output ke numpy
+        demand_pred = demand_pred.cpu().numpy().flatten().tolist()
+        
+        # ========== [2] CHRONOS-T5 DECISION MAKING ==========
+        chronos = pipeline(
+            "text-generation",
+            model="amazon/chronos-t5-small",
+            device=device
+        )
+        
+        prompt = f"""
+        Data UMKM:
+        - Prediksi demand 7 hari ke depan: {demand_pred}
+        - Stok saat ini: {data['supply'].iloc[-1]}
+        - Harga rata-rata: {data['harga'].mean() if 'harga' in data else 'N/A'}
+        Berikan rekomendasi pembelian dalam 1 kalimat:
+        """
+        
+        recommendation = chronos(prompt, max_length=100)[0]['generated_text']
+        return recommendation
     
-    # Prediksi demand
-    inputs = {"values": data['demand'].tolist()}
-    demand_pred = ttm.generate(**inputs, max_length=7)
-    
-    # Rekomendasi dengan Chronos-T5
-    chronos = pipeline(
-        "text-generation", 
-        model="amazon/chronos-t5-small",
-        device=device
-    )
-    prompt = f"""
-    Data UMKM:
-    - Prediksi demand: {demand_pred}
-    - Stok saat ini: {data['supply'].iloc[-1]}
-    Beri rekomendasi dalam 1 kalimat:
-    """
-    return chronos(prompt, max_length=50)[0]['generated_text']
+    except Exception as e:
+        return f"Error: {str(e)}"