app.py

# app.py
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import json
from jsonschema import validate, ValidationError
import logging

# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("StrategyInterpreterSpace")

# Load model and tokenizer
model_name = "EleutherAI/gpt-neo-2.7B"  # Updated model
logger.info(f"Loading model '{model_name}'...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
logger.info("Model loaded successfully.")

# Define JSON schema
schema = {
    "type": "object",
    "required": [
        "strategy_name",
        "market_type",
        "assets",
        "trade_parameters",
        "conditions",
        "risk_management"
    ],
    "properties": {
        "strategy_name": {"type": "string"},
        "market_type": {"type": "string", "enum": ["spot", "futures", "margin"]},
        "assets": {"type": "array", "items": {"type": "string"}},
        "trade_parameters": {
            "type": "object",
            "required": ["leverage", "order_type", "position_size"],
            "properties": {
                "leverage": {"type": "number"},
                "order_type": {"type": "string"},
                "position_size": {"type": "number"}
            }
        },
        "conditions": {
            "type": "object",
            "required": ["entry", "exit"],
            "properties": {
                "entry": {
                    "type": "array",
                    "items": {"$ref": "#/definitions/condition"}
                },
                "exit": {
                    "type": "array",
                    "items": {"$ref": "#/definitions/condition"}
                }
            }
        },
        "risk_management": {
            "type": "object",
            "required": ["stop_loss", "take_profit", "trailing_stop_loss"],
            "properties": {
                "stop_loss": {"type": "number"},
                "take_profit": {"type": "number"},
                "trailing_stop_loss": {"type": "number"}
            }
        }
    },
    "definitions": {
        "condition": {
            "type": "object",
            "required": ["indicator", "operator", "value", "timeframe"],
            "properties": {
                "indicator": {"type": "string"},
                "operator": {"type": "string", "enum": [">", "<", "==", ">=", "<="]},
                "value": {"type": ["string", "number"]},
                "timeframe": {"type": "string"},
                "indicator_parameters": {
                    "type": "object",
                    "properties": {
                        "period": {"type": "number"},
                    },
                    "additionalProperties": True
                }
            }
        }
    }
}

def interpret_strategy(description: str) -> str:
    prompt = f"""
You are an expert crypto trading assistant. Convert the following trading strategy description into a JSON format 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. 
Strategy Description:
{description}

JSON:
"""
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=1000,
        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()

    # 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:
    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)}

Use indicators and conditions that can be applied by ccxt, bitget, pandas-ta, and backtrader.

JSON:"""

    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(
        inputs,
        max_length=1000,
        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()

    # Validate JSON
    try:
        strategy_data = json.loads(response_text)
        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}"

iface_interpret = gr.Interface(
    fn=interpret_strategy,
    inputs=gr.inputs.Textbox(lines=10, placeholder="Enter your strategy description here..."),
    outputs="text",
    title="Strategy Interpreter",
    description="Convert trading strategy descriptions into structured JSON format."
)

iface_suggest = gr.Interface(
    fn=suggest_strategy,
    inputs=[
        gr.inputs.Textbox(lines=1, placeholder="Enter risk level (e.g., medium)...", label="Risk Level"),
        gr.inputs.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()