Update app.py

app.py

@@ -4,17 +4,22 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import json
 from jsonschema import validate, ValidationError
 import logging
+import torch
 
 # Initialize logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("StrategyInterpreterSpace")
 
 # Load model and tokenizer
-model_name = "EleutherAI/gpt-neo-2.7B"  # Updated model
+model_name = "EleutherAI/gpt-neo-2.7B"  # Using a smaller model to fit within memory constraints
 logger.info(f"Loading model '{model_name}'...")
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-logger.info("Model loaded successfully.")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    logger.info("Model loaded successfully.")
+except Exception as e:
+    logger.error(f"Failed to load model: {e}")
+    raise e
 
 # Define JSON schema
 schema = {
@@ -85,41 +90,69 @@ schema = {
     }
 }
 
-def interpret_strategy(description: str) -> str:
+def interpret_strategy(description: str) -> dict:
+    logger.info("Received strategy description for interpretation.")
     prompt = f"""
-You are an expert crypto trading assistant. Convert the following trading strategy description into a JSON format following this schema:
+You are an expert crypto trading assistant. Convert the following trading strategy description into a JSON format strictly following this schema:
 
 {json.dumps(schema, indent=2)}
 
-Include all indicators (only ones available in Ta-lib and pandas-ta), their parameters (only ones that are standard for ccxt and backtrader to support), assets (only ones that are available through BitGet) as trading pairs, conditions (only those supported by bitget, backtrader, finta, pandas-ta), risk management settings, and trade execution details (only those supported by ccxt, bitget and backtrader). 
-Response should only return the JSON with the correct parameters, nothing else. 
+Ensure that the response contains only valid JSON with the correct parameters. Do not include any additional text or explanations.
+
 Strategy Description:
 {description}
 
 JSON:
 """
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(
-        inputs,
-        max_length=1000,  # You can adjust this value as needed
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response_text = generated_text[len(prompt):].strip()
+    logger.debug(f"Prompt constructed: {prompt}")
+
+    try:
+        inputs = tokenizer.encode(prompt, return_tensors="pt")
+        logger.info("Tokenized the input prompt.")
+    except Exception as e:
+        logger.error(f"Error during tokenization: {e}")
+        return {"error": f"Error during tokenization: {e}"}
+
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=800,  # Reduced max_length to prevent overly long outputs
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        logger.info("Model generated output.")
+    except Exception as e:
+        logger.error(f"Error during model generation: {e}")
+        return {"error": f"Error during model generation: {e}"}
+
+    try:
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logger.debug(f"Generated text: {generated_text}")
+        response_text = generated_text[len(prompt):].strip()
+        logger.debug(f"Response text after prompt removal: {response_text}")
+    except Exception as e:
+        logger.error(f"Error during decoding: {e}")
+        return {"error": f"Error during decoding: {e}"}
 
     # Validate JSON
     try:
         strategy_data = json.loads(response_text)
         validate(instance=strategy_data, schema=schema)
-        return json.dumps(strategy_data, indent=2)
-    except (json.JSONDecodeError, ValidationError) as e:
-        logger.error(f"Error interpreting strategy: {e}")
-        return f"Error interpreting strategy: {e}"
-
-def suggest_strategy(risk_level: str, market_type: str) -> str:
+        logger.info("Strategy interpreted successfully and validated against schema.")
+        return strategy_data  # Return as dict for Gradio's JSON output
+    except json.JSONDecodeError as e:
+        logger.error(f"JSON decoding error: {e}")
+        # Return raw text for debugging
+        return {"error": f"Error interpreting strategy: Invalid JSON format.\nGenerated Text:\n{response_text}\nDetails: {e}"}
+    except ValidationError as e:
+        logger.error(f"JSON validation error: {e}")
+        return {"error": f"Error interpreting strategy: JSON does not conform to schema.\nDetails: {e}"}
+
+def suggest_strategy(risk_level: str, market_type: str) -> dict:
+    logger.info("Received request to suggest a new strategy.")
     prompt = f"""Please create a unique crypto trading strategy suitable for a '{risk_level}' risk appetite in the '{market_type}' market.
 Ensure the JSON matches this schema:
 {json.dumps(schema, indent=2)}
@@ -127,18 +160,38 @@ Ensure the JSON matches this schema:
 Use indicators and conditions that can be applied by ccxt, bitget, pandas-ta, and backtrader.
 
 JSON:"""
+    logger.debug(f"Prompt constructed for strategy suggestion: {prompt}")
 
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
-    outputs = model.generate(
-        inputs,
-        max_length=1000,  # You can adjust this value as needed
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    response_text = generated_text[len(prompt):].strip()
+    try:
+        inputs = tokenizer.encode(prompt, return_tensors="pt")
+        logger.info("Tokenized the suggestion prompt.")
+    except Exception as e:
+        logger.error(f"Error during tokenization: {e}")
+        return {"error": f"Error during tokenization: {e}"}
+
+    try:
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs,
+                max_length=800,  # Reduced max_length to prevent overly long outputs
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+        logger.info("Model generated suggestion output.")
+    except Exception as e:
+        logger.error(f"Error during model generation: {e}")
+        return {"error": f"Error during model generation: {e}"}
+
+    try:
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logger.debug(f"Generated suggestion text: {generated_text}")
+        response_text = generated_text[len(prompt):].strip()
+        logger.debug(f"Suggestion response text after prompt removal: {response_text}")
+    except Exception as e:
+        logger.error(f"Error during decoding: {e}")
+        return {"error": f"Error during decoding: {e}"}
 
     # Validate JSON
     try:
@@ -146,15 +199,23 @@ JSON:"""
         validate(instance=strategy_data, schema=schema)
         if strategy_data.get("market_type") != market_type:
             raise ValueError("The generated strategy's market type does not match the selected market type.")
-        return json.dumps(strategy_data, indent=2)
-    except (json.JSONDecodeError, ValidationError, ValueError) as e:
-        logger.error(f"Error generating strategy: {e}")
-        return f"Error generating strategy: {e}"
+        logger.info("Strategy suggested successfully and validated against schema.")
+        return strategy_data  # Return as dict for Gradio's JSON output
+    except json.JSONDecodeError as e:
+        logger.error(f"JSON decoding error: {e}")
+        # Return raw text for debugging
+        return {"error": f"Error generating strategy: Invalid JSON format.\nGenerated Text:\n{response_text}\nDetails: {e}"}
+    except ValidationError as e:
+        logger.error(f"JSON validation error: {e}")
+        return {"error": f"Error generating strategy: JSON does not conform to schema.\nDetails: {e}"}
+    except ValueError as e:
+        logger.error(f"Market type mismatch error: {e}")
+        return {"error": f"Error generating strategy: {e}"}
 
 iface_interpret = gr.Interface(
     fn=interpret_strategy,
-    inputs=gr.Textbox(lines=10, placeholder="Enter your strategy description here..."),
-    outputs="text",
+    inputs=gr.Textbox(lines=10, placeholder="Enter your strategy description here...", label="Strategy Description"),
+    outputs=gr.JSON(label="Interpreted Strategy"),
     title="Strategy Interpreter",
     description="Convert trading strategy descriptions into structured JSON format."
 )
@@ -163,14 +224,4 @@ iface_suggest = gr.Interface(
     fn=suggest_strategy,
     inputs=[
         gr.Textbox(lines=1, placeholder="Enter risk level (e.g., medium)...", label="Risk Level"),
-        gr.Textbox(lines=1, placeholder="Enter market type (e.g., spot)...", label="Market Type")
-    ],
-    outputs="text",
-    title="Strategy Suggester",
-    description="Generate a unique trading strategy based on risk level and market type."
-)
-
-app = gr.TabbedInterface([iface_interpret, iface_suggest], ["Interpret Strategy", "Suggest Strategy"])
-
-if __name__ == "__main__":
-    app.launch()
+