Update utils/model.py

utils/model.py

@@ -1,32 +1,47 @@
 from transformers import pipeline, AutoModelForTimeSeriesPrediction
 import torch
 import numpy as np
+import pandas as pd
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def predict_umkm(data):
     try:
-        # ===== 1. GRANITE-TTM =====
-        # Pastikan kolom 'demand' ada dan tidak kosong
-        if 'demand' not in data or len(data['demand']) == 0:
-            raise ValueError("Kolom 'demand' tidak ditemukan atau kosong")
+        # ===== 1. Validasi Input =====
+        if not isinstance(data, pd.DataFrame):
+            data = pd.DataFrame(data)
+            
+        required_cols = ['demand', 'supply']
+        for col in required_cols:
+            if col not in data.columns:
+                raise ValueError(f"Kolom {col} tidak ditemukan")
         
-        # Konversi data ke format yang diterima GRANITE
-        values = torch.tensor(data['demand'].values, dtype=torch.float32).unsqueeze(0).to(device)
+        # ===== 2. GRANITE-TTM Forecasting =====
+        # Konversi data ke tensor dengan format khusus
+        values = torch.tensor(
+            data['demand'].astype(float).values,
+            dtype=torch.float32
+        ).unsqueeze(0).unsqueeze(-1).to(device)  # Shape: [1, seq_len, 1]
         
-        # Load model
+        # Load model (pastikan nama model benar)
         ttm_model = AutoModelForTimeSeriesPrediction.from_pretrained(
-            "ibm/granite-timeseries-ttm-r2"
+            "ibm/granite-ttm-r2",
+            trust_remote_code=True
         ).to(device)
         
-        # Prediksi (pastikan input shape [batch_size, sequence_length])
+        # Generate predictions
         with torch.no_grad():
-            predictions = ttm_model.generate(input_values=values, max_length=7)
+            predictions = ttm_model.generate(
+                inputs=values,
+                max_length=min(7, len(data)+3)  # Prediksi maks 7 hari
+            )
         
-        # Konversi output ke list
-        demand_pred = predictions.cpu().numpy().flatten().tolist()
+        # Post-processing output
+        demand_pred = predictions.cpu().numpy().squeeze().tolist()
+        if isinstance(demand_pred, float):
+            demand_pred = [demand_pred]  # Convert single value to list
         
-        # ===== 2. CHRONOS-T5 =====
+        # ===== 3. Chronos-T5 Decision =====
         chronos = pipeline(
             "text-generation",
             model="amazon/chronos-t5-small",
@@ -34,13 +49,27 @@ def predict_umkm(data):
         )
         
         prompt = f"""
-        Data UMKM:
-        - Prediksi 7 hari: {demand_pred}
-        - Stok saat ini: {data['supply'].iloc[-1]}
-        Berikan rekomendasi dalam 1 kalimat:
+        [INSTRUCTION]
+        Berikan rekomendasi inventory untuk UMKM berdasarkan:
+        - Prediksi demand 7 hari: {demand_pred[:7]}
+        - Stok saat ini: {data['supply'].iloc[-1]:.0f}
+        - Trend: {'naik' if demand_pred[-1] > demand_pred[0] else 'turun'}
+        
+        [FORMAT]
+        - Gunakan bahasa Indonesia
+        - Maksimal 1 kalimat
+        - Sertakan angka konkret
+        [/FORMAT]
         """
         
-        return chronos(prompt, max_length=100)[0]['generated_text']
+        recommendation = chronos(
+            prompt,
+            max_new_tokens=50,
+            do_sample=True,
+            temperature=0.7
+        )[0]['generated_text']
+        
+        return recommendation.split("[/FORMAT]")[-1].strip()
     
     except Exception as e:
-        return f"⚠️ Error: {str(e)}"
+        return f"🚨 Error: {str(e)}"