Spaces:

mike23415
/

Reserch

Runtime error

App Files Files Community

mike23415 commited on 19 days ago

Commit

9d13a5a

verified ·

1 Parent(s): 1b26681

Update app.py

Browse files

Files changed (1) hide show

app.py +109 -53

app.py CHANGED Viewed

@@ -1,37 +1,63 @@
 from flask import Flask, request, send_file, jsonify
 from flask_cors import CORS
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import pdfkit
-import jinja2
-import torch
-import tempfile
-import os
-os.environ['TRANSFORMERS_CACHE'] = '/app/.cache'
-os.environ['HF_DATASETS_CACHE'] = '/app/.cache'
-os.environ['XDG_CACHE_HOME'] = '/app/.cache'
 os.environ['HF_HOME'] = '/app/.cache'
 app = Flask(__name__)
 CORS(app)
-# Initialize model and tokenizer
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        "gpt2-medium",
-        from_tf=False,
-        use_safetensors=True
-    )
-    tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
-    generator = pipeline(
-        'text-generation',
-        model=model,
-        tokenizer=tokenizer,
-        device=0 if torch.cuda.is_available() else -1
-    )
-except Exception as e:
-    print(f"Model loading failed: {str(e)}")
-    generator = None
 IEEE_TEMPLATE = """
 <!DOCTYPE html>
 <html>
@@ -45,14 +71,8 @@ IEEE_TEMPLATE = """
             font-size: 12pt;
             line-height: 1.5;
         }
-        .header {
-            text-align: center;
-            margin-bottom: 24pt;
-        }
-        .two-column {
-            column-count: 2;
-            column-gap: 0.5in;
-        }
         h1 { font-size: 14pt; margin: 12pt 0; }
         h2 { font-size: 12pt; margin: 12pt 0 6pt 0; }
         .abstract { margin-bottom: 24pt; }
@@ -97,31 +117,46 @@ IEEE_TEMPLATE = """
 </html>
 """
-def format_content(content):
-    if not generator:
-        return content  # Fallback if model failed to load
-    try:
-        prompt = f"Format this research content to IEEE standards:\n{str(content)}"
-        return generator(
-            prompt,
-            max_length=1024,
-            num_return_sequences=1,
-            clean_up_tokenization_spaces=True
-        )[0]['generated_text']
-    except Exception as e:
-        print(f"Formatting failed: {str(e)}")
-        return content
 @app.route('/generate', methods=['POST'])
 def generate_pdf():
     try:
         data = request.json
-        if not data or 'title' not in data or 'authors' not in data:
-            return jsonify({"error": "Missing required fields"}), 400
-        # Format content using AI
-        formatted = format_content(data.get('content', {}))
         # Generate HTML
         html = jinja2.Template(IEEE_TEMPLATE).render(
@@ -144,13 +179,34 @@ def generate_pdf():
             'quiet': ''
         }
-        # Create PDF
         with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as f:
             pdfkit.from_string(html, f.name, options=options)
             return send_file(f.name, mimetype='application/pdf')
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)

+import os
+import time
+import tempfile
+import jinja2
+import pdfkit
+import torch
+from threading import Thread
 from flask import Flask, request, send_file, jsonify
 from flask_cors import CORS
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+# Configure cache directories
 os.environ['HF_HOME'] = '/app/.cache'
+os.environ['XDG_CACHE_HOME'] = '/app/.cache'
 app = Flask(__name__)
 CORS(app)
+# Global state tracking
+model_loaded = False
+load_error = None
+generator = None
+# --------------------------------------------------
+# Asynchronous Model Loading
+# --------------------------------------------------
+def load_model():
+    global model_loaded, load_error, generator
+    try:
+        # Initialize model with low-memory settings
+        model = AutoModelForCausalLM.from_pretrained(
+            "gpt2-medium",
+            use_safetensors=True,
+            device_map="auto",
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        )
+        tokenizer = AutoTokenizer.from_pretrained("gpt2-medium")
+        generator = pipeline(
+            'text-generation',
+            model=model,
+            tokenizer=tokenizer,
+            device=0 if torch.cuda.is_available() else -1
+        )
+        model_loaded = True
+        print("Model loaded successfully")
+    except Exception as e:
+        load_error = str(e)
+        print(f"Model loading failed: {load_error}")
+# Start model loading in background thread
+Thread(target=load_model).start()
+# --------------------------------------------------
+# IEEE Format Template
+# --------------------------------------------------
 IEEE_TEMPLATE = """
 <!DOCTYPE html>
 <html>
             font-size: 12pt;
             line-height: 1.5;
         }
+        .header { text-align: center; margin-bottom: 24pt; }
+        .two-column { column-count: 2; column-gap: 0.5in; }
         h1 { font-size: 14pt; margin: 12pt 0; }
         h2 { font-size: 12pt; margin: 12pt 0 6pt 0; }
         .abstract { margin-bottom: 24pt; }
 </html>
 """
+# --------------------------------------------------
+# API Endpoints
+# --------------------------------------------------
+@app.route('/health', methods=['GET'])
+def health_check():
+    if load_error:
+        return jsonify({
+            "status": "error",
+            "message": f"Model failed to load: {load_error}"
+        }), 500
+    return jsonify({
+        "status": "ready" if model_loaded else "loading",
+        "model_loaded": model_loaded,
+        "device": "cuda" if torch.cuda.is_available() else "cpu"
+    }), 200 if model_loaded else 503
 @app.route('/generate', methods=['POST'])
 def generate_pdf():
+    # Check model status
+    if not model_loaded:
+        return jsonify({
+            "error": "Model not loaded yet",
+            "status": "loading"
+        }), 503
     try:
+        # Validate input
         data = request.json
+        if not data:
+            return jsonify({"error": "No data provided"}), 400
+        required = ['title', 'authors', 'content']
+        if missing := [field for field in required if field not in data]:
+            return jsonify({
+                "error": f"Missing fields: {', '.join(missing)}"
+            }), 400
+        # Format content
+        formatted = format_content(data['content'])
         # Generate HTML
         html = jinja2.Template(IEEE_TEMPLATE).render(
             'quiet': ''
         }
+        # Create temporary PDF
         with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as f:
             pdfkit.from_string(html, f.name, options=options)
             return send_file(f.name, mimetype='application/pdf')
     except Exception as e:
         return jsonify({"error": str(e)}), 500
+    finally:
+        if 'f' in locals():
+            try: os.remove(f.name)
+            except: pass
+# --------------------------------------------------
+# Content Formatting
+# --------------------------------------------------
+def format_content(content):
+    try:
+        prompt = f"Format this research content to IEEE standards:\n{str(content)}"
+        return generator(
+            prompt,
+            max_new_tokens=512,
+            temperature=0.7,
+            do_sample=True,
+            truncation=True
+        )[0]['generated_text']
+    except Exception as e:
+        print(f"Formatting error: {str(e)}")
+        return content
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)