app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import json

MODEL_NAME = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

def set_token(json_input):
    """
    Accepts JSON (string or dict) and returns a JSON response.
    """
    try:
        if isinstance(json_input, str):
            data = json.loads(json_input)
        else:
            data = json_input
    except Exception as e:
        return {"error": f"Invalid JSON input: {str(e)}"}

    system_prompt = data.get("system_prompt", "You are a helpful AI.")
    user_input = data.get("user_input", "")
    temperature = float(data.get("temperature", 0.7))
    max_tokens = int(data.get("max_tokens", 100))

    prompt = f"{system_prompt}\nUser: {user_input}\nAI:"
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        pad_token_id=tokenizer.eos_token_id,
    )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    if "AI:" in text:
        text = text.split("AI:")[-1].strip()

    return {
        "model": MODEL_NAME,
        "response": text,
        "tokens_used": len(outputs[0]),
    }

# Gradio interface
iface = gr.Interface(
    fn=set_token,
    inputs=gr.Textbox(label="JSON Request", placeholder='{"system_prompt":"...","user_input":"..."}'),
    outputs="json",
    title="setToken API",
    description="Send a JSON request with system_prompt and user_input. Example: {\"system_prompt\":\"You are helpful.\",\"user_input\":\"Say hi!\"}"
)

iface.launch()