Update app.py

app.py

@@ -2,15 +2,27 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer
 from model_smol2 import LlamaForCausalLM, config_model
-
+import requests
 # Instantiate the model
 model = LlamaForCausalLM(config_model)
 
-# Load the checkpoint
-checkpoint_path = "/Users/shriti/Downloads/Assign13_ERAV3/deply/final_checkpoint.pt"
-checkpoint = torch.load(checkpoint_path, map_location="cpu")
-model.load_state_dict(checkpoint['model_state_dict'])
-model.eval()
+# Correct URL for Google Drive direct download
+url = "https://drive.google.com/uc?id=1tyZhudOcZRMaSLkzAWksVyQmbPQ0fi50"
+
+response = requests.get(url)
+if response.status_code == 200:
+    with open("final_checkpoint.pt", "wb") as f:
+        f.write(response.content)
+else:
+    print(f"Failed to download the file. Status code: {response.status_code}")
+
+# Now load the checkpoint
+try:
+    checkpoint = torch.load("final_checkpoint.pt", map_location="cpu", weights_only=True)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+except Exception as e:
+    print(f"Error loading checkpoint: {e}")
 
 # Load tokenizer (replace with the appropriate tokenizer if you're using a custom one)
 # Load the tokenizer