Upload folder using huggingface_hub

app.py

@@ -4,20 +4,30 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from peft import PeftModel
 import torch
 
+# Determine device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained("rxpbtn21/t5-small-lora-summarizer")
 
 # Load the base model and then the LoRA adapter
-base_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small", device_map="auto")
+# Ensure the base model is also moved to the correct device
+base_model = AutoModelForSeq2SeqLM.from_pretrained("t5-small").to(device)
 model = PeftModel.from_pretrained(base_model, "rxpbtn21/t5-small-lora-summarizer")
 model.eval()
 
 def summarize(text):
-    inputs = tokenizer(text, max_length=512, truncation=True, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(inputs["input_ids"].to(model.device), num_beams=4, max_new_tokens=128, early_stopping=True)
-    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return summary
+    try:
+        inputs = tokenizer(text, max_length=512, truncation=True, return_tensors="pt")
+        with torch.no_grad():
+            # Ensure inputs are on the same device as the model
+            outputs = model.generate(inputs["input_ids"].to(device), num_beams=4, max_new_tokens=128, early_stopping=True)
+        summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return summary
+    except Exception as e:
+        # Log the error and return an informative message
+        print(f"Error during summarization: {e}")
+        return f"An error occurred during summarization. Please check the Space logs for details. Error: {e}"
 
 # Create Gradio interface
 iface = gr.Interface(fn=summarize, inputs="text", outputs="text", title="LoRA Fine-tuned T5-small Summarizer")