Update app.py

app.py

@@ -10,12 +10,18 @@ from peft import LoraConfig, get_peft_model
 from datasets import Dataset
 import os
 from transformers import Trainer, TrainingArguments
-from huggingface_hub import login, upload_file
+from huggingface_hub import login, upload_file, hf_hub_download
 import wandb
 import gradio as gr
 import spaces
+from huggingface_hub import HfApi
+
+hf_api = HfApi()
+
 HF_TOKEN = os.getenv('HF_TOKEN')
 WANDB_TOKEN = os.getenv('WANDB_TOKEN')
+login(token=HF_TOKEN)
+wandb.login(key=WANDB_TOKEN)
 
 class Scaler:
     def __init__(self, feature_range):
@@ -38,14 +44,24 @@ class Scaler:
         min_x, max_x = np.min(X), np.max(X)
         return (X - min_x) / (max_x - min_x) * (max_val - min_val) + min_val
 
+def check_existing_model(stock_symbol, start_date, end_date):
+    repo_id = f"Q-bert/StockLlama-tuned-{stock_symbol}-{stock_symbol}-{start_date}_{end_date}"
+    state = repo_id in [model.modelId for model in hf_api.list_models()]
+    return state
+    
 @spaces.GPU
 def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100), data_seq_length=256, epochs=10, batch_size=16, learning_rate=2e-4):
+    repo_id = f"Q-bert/StockLlama-tuned-{stock_symbol}-{stock_symbol}-{start_date}_{end_date}"
+
+    if check_existing_model(stock_symbol, start_date, end_date):
+        return f"Model for {stock_symbol} from {start_date} to {end_date} already exists."
+
     try:
         stock_data = yf.download(stock_symbol, start=start_date, end=end_date, progress=False)
     except Exception as e:
         print(f"Error downloading data for {stock_symbol}: {e}")
         return
-    
+
     data = stock_data["Close"]
 
     scaler = Scaler(feature_range)
@@ -55,7 +71,7 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
     seq = [np.array(scaled_data[i:i + data_seq_length]) for i in range(len(scaled_data) - data_seq_length)]
     target = [np.array(scaled_data[i + data_seq_length:i + data_seq_length + 1]) for i in range(len(scaled_data) - data_seq_length)]
 
-    seq_tensors = [torch.tensor(s, dtype=torch.float32).unsqueeze(0) for s in seq]
+    seq_tensors = [torch.tensor(s, dtype=torch.float32) for s in seq]
     target_tensors = [t[0] for t in target]
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -70,11 +86,9 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
     )
     model = get_peft_model(model, config)
 
-    login(token=HF_TOKEN)
-    wandb.login(key=WANDB_TOKEN)
     dct = {"input_ids": seq_tensors, "label": target_tensors}
     dataset = Dataset.from_dict(dct)
-    dataset.push_to_hub(f"Q-bert/{stock_symbol}-{start_date}_{end_date}")
+    dataset.push_to_hub(repo_id)
     trainer = Trainer(
         model=model,
         train_dataset=dataset,
@@ -108,11 +122,12 @@ def train_stock_model(stock_symbol, start_date, end_date, feature_range=(10, 100
         path_in_repo=f"scalers/{scaler_path}",
         repo_id=f"Q-bert/StockLlama-tuned-{stock_symbol}-{stock_symbol}-{start_date}_{end_date}"
     )
+    return f"Training completed and model saved for {stock_symbol} from {start_date} to {end_date}."
 
 @spaces.GPU
 def gradio_train_stock_model(stock_symbol, start_date, end_date, feature_range_min, feature_range_max, data_seq_length, epochs, batch_size, learning_rate):
     feature_range = (feature_range_min, feature_range_max)
-    train_stock_model(
+    result = train_stock_model(
         stock_symbol=stock_symbol,
         start_date=start_date,
         end_date=end_date,
@@ -122,7 +137,7 @@ def gradio_train_stock_model(stock_symbol, start_date, end_date, feature_range_m
         batch_size=batch_size,
         learning_rate=learning_rate
     )
-    return f"Training initiated for {stock_symbol} from {start_date} to {end_date}."
+    return result
 
 iface = gr.Interface(
     fn=gradio_train_stock_model,
@@ -140,4 +155,4 @@ iface = gr.Interface(
     outputs="text",
 )
 
-iface.launch()
+iface.launch()