Update app.py

app.py

@@ -2,6 +2,7 @@ from flask import Flask, request, jsonify
 import torch
 from transformers import pipeline
 import os
+import json # <-- Humne yeh nayi library import ki hai
 
 # Model ka naam
 model_name = "distilgpt2"
@@ -12,8 +13,8 @@ try:
     pipe = pipeline(
         "text-generation",
         model=model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="auto"
+        torch_dtype=torch.float32, # CPU ke liye float32 behtar hai
+        device=-1 # CPU par run karne ke liye force karna
     )
     print("✅ Model loaded successfully!")
 except Exception as e:
@@ -42,24 +43,30 @@ def generate_content():
     data = request.get_json()
     if not data or 'step_text' not in data:
         return jsonify({"error": "step_text is missing from the request body"}), 400
-
+    
     step_text = data.get('step_text')
 
     try:
         final_prompt = meta_prompt_template.format(step_text=step_text)
-
+        
         outputs = pipe(final_prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
-
+        
         generated_text = outputs[0]['generated_text']
-
-        # Model ke output se saaf JSON nikalna
-        json_part = generated_text[generated_text.find('{'):generated_text.rfind('}')+1]
-
-        # String ko JSON mein convert karna
-        response_data = eval(json_part)
-
-        return jsonify(response_data)
-
+        
+        # Naya, behtar tareeka JSON nikalne ka
+        start_index = generated_text.find('{')
+        end_index = generated_text.rfind('}') + 1
+        
+        if start_index != -1 and end_index != -1:
+            json_part = generated_text[start_index:end_index]
+            # String ko JSON mein convert karna
+            response_data = json.loads(json_part)
+            return jsonify(response_data)
+        else:
+            return jsonify({"error": "Failed to find a valid JSON object in the model's response.", "raw_response": generated_text}), 500
+
+    except json.JSONDecodeError:
+        return jsonify({"error": "Model returned a malformed JSON string.", "raw_response": generated_text}), 500
     except Exception as e:
         print(f"Error during generation: {e}")
         return jsonify({"error": str(e)}), 500