Update app.py

app.py

@@ -1,35 +1,31 @@
-from flask import Flask, request, jsonify
-from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
+import gradio as gr
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
+import json
 
-app = Flask(__name__)
-
-# Load a lightweight model for structured output
-model_name = "google/flan-t5-small"  # You can replace this with another small instruct model
+# Load lightweight model
+model_name = "google/flan-t5-small"  # You can change this to another small instruct model
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
 generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-@app.route('/generate-json', methods=['POST'])
-def generate_json():
-    data = request.json
-    prompt = data.get("prompt")
-
-    if not prompt:
-        return jsonify({"error": "Prompt is required"}), 400
-
-    # Add instruction if needed
+def generate_json(prompt):
     instruction = f"Generate a JSON object from the following description:\n{prompt}"
-    
-    response = generator(instruction, max_length=256, do_sample=False)
-    generated_text = response[0]["generated_text"]
+    result = generator(instruction, max_length=256, do_sample=False)
+    generated_text = result[0]["generated_text"]
 
     try:
-        # Try to parse generated text as JSON
-        generated_json = eval(generated_text)  # Use json.loads if safe JSON string is expected
-        return jsonify({"json": generated_json})
+        parsed = eval(generated_text)  # Not recommended in production; use json.loads if output is valid JSON
+        formatted_json = json.dumps(parsed, indent=2)
     except Exception as e:
-        return jsonify({"raw_output": generated_text, "error": str(e)}), 200
+        formatted_json = f"Raw Output:\n{generated_text}\n\nError parsing JSON: {e}"
+
+    return formatted_json
 
-if __name__ == '__main__':
-    app.run(debug=True)
+gr.Interface(
+    fn=generate_json,
+    inputs=gr.Textbox(lines=4, label="Enter Prompt"),
+    outputs=gr.Textbox(lines=20, label="Generated JSON"),
+    title="Lightweight JSON Generator",
+    description="Enter a prompt describing the structure or content you want in JSON format."
+).launch()