Update app.py

app.py

@@ -1,55 +1,55 @@
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import json
 
-app = Flask(__name__)
-
-# Load your model (can be any HF model)
-MODEL_NAME = "tiiuae/falcon-7b-instruct"
+MODEL_NAME = "gpt2"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
-)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
 
-@app.route('/setToken', methods=['POST'])
-def set_token():
+def set_token(json_input):
     """
-    Main multimodal API endpoint.
-    Handles system + user prompts and returns generated response.
+    Accepts JSON (string or dict) and returns a JSON response.
     """
-    data = request.get_json(force=True)
+    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))
-    mode = data.get("mode", "text")
-
-    # Text mode (default)
-    if mode == "text":
-        full_prompt = f"{system_prompt}\nUser: {user_input}\nAI:"
-        inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=temperature
-        )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = response.split("AI:")[-1].strip()
-        return jsonify({
-            "model": MODEL_NAME,
-            "response": response,
-            "mode": "text"
-        })
-
-    # You can later add multimodal branches here:
-    # - "image" -> call image generation pipeline
-    # - "audio" -> call speech-to-text / text-to-speech
-    # - "embedding" -> return vector embeddings
-
-    return jsonify({"error": f"Unsupported mode: {mode}"}), 400
+    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!\"}"
+)
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+iface.launch()