Upload folder using huggingface_hub

app.py

@@ -1,15 +1,17 @@
 """
-Gradio Space: Python docstring generation. Run: gradio app.py
+Gradio Space: Python docstring generation.
+HF runs the app; do not call demo.launch() here.
 """
 import gradio as gr
-from inference import generate_docstring
-
 
 def summarize_code(code: str) -> str:
     if not code or not code.strip():
         return "Paste a Python code snippet above."
-    return generate_docstring(code, model_name="t5-small", max_length=128, num_beams=4)
-
+    try:
+        from inference import generate_docstring
+        return generate_docstring(code, model_name="t5-small", max_length=128, num_beams=4)
+    except Exception as e:
+        return f"Error: {str(e)}. (Model may be loading; try again.)"
 
 demo = gr.Interface(
     fn=summarize_code,
@@ -22,5 +24,3 @@ demo = gr.Interface(
     title="Python Docstring Generator",
     description="Paste a Python function or snippet to get a short docstring summary.",
 )
-
-demo.launch()

inference.py

@@ -1,9 +1,9 @@
 """
-Inference for docstring generation. Uses T5.
+Inference for docstring generation. Uses T5 (cached after first load).
 """
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 
+_cache = {}
 
 def generate_docstring(
     code: str,
@@ -14,8 +14,14 @@ def generate_docstring(
 ) -> str:
     if device is None:
         device = "cuda" if torch.cuda.is_available() else "cpu"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+    if model_name not in _cache:
+        from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+        _cache[model_name] = {
+            "tokenizer": AutoTokenizer.from_pretrained(model_name),
+            "model": AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device),
+        }
+    tokenizer = _cache[model_name]["tokenizer"]
+    model = _cache[model_name]["model"]
     input_text = "summarize: " + code
     inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512).to(device)
     with torch.no_grad():

requirements.txt

@@ -1,3 +1,4 @@
 torch>=2.0.0
 transformers>=4.30.0
-gradio>=4.0.0
+gradio>=4.0.0,<5.0.0
+accelerate>=0.20.0