# Install necessary packages (only for local environment)
# !pip install pandas gradio transformers torch

import pandas as pd
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the Hugging Face forecasting model
def load_model():
    model_name = "Ankur87/Llama2_Time_series_forecasting_7.0"  # Using the specified model
    model = AutoModelForCausalLM.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = load_model()

def forecast(csv_file):
    data = pd.read_csv(csv_file.name, parse_dates=['timestamp_column'])
    
    # Convert data to a format suitable for the model
    input_text = data.to_json()
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
    
    with torch.no_grad():
        predictions = model.generate(**inputs, max_length=200)
    
    forecast_text = tokenizer.decode(predictions[0], skip_special_tokens=True)
    forecasts = pd.DataFrame({'forecast': [forecast_text]})
    forecasts.to_csv("forecasts.csv", index=False)
    return "forecasts.csv"

# Gradio Interface
iface = gr.Interface(
    fn=forecast,
    inputs=gr.File(label="Upload CSV File"),
    outputs=gr.File(label="Download Forecasts"),
    title="Time Series Forecasting with Llama2",
    description="Upload a CSV file with a timestamp column to generate forecasts using Llama2."
)

iface.launch()