Update app.py

app.py

@@ -1,26 +1,13 @@
-from flask import Flask, request, jsonify
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+import gradio as gr
 
 # Initialize the model and tokenizer
 model = AutoModelForCausalLM.from_pretrained("anto18671/lumenspark", trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("anto18671/lumenspark", trust_remote_code=True)
 
-# Set up Flask application
-app = Flask(__name__)
-
-# Define inference endpoint
-@app.route("/generate", methods=["POST"])
-def generate_text():
-    data = request.get_json()
-
-    # Extract the input text
-    text = data.get("text", "")
-    if not text:
-        return jsonify({"error": "Input text is required"}), 400
-
+def generate_text(input_text):
     # Tokenize input text
-    encoded_input = tokenizer(text, return_tensors='pt')
+    encoded_input = tokenizer(input_text, return_tensors='pt')
 
     # Generate text using the model
     output = model.generate(
@@ -37,10 +24,16 @@ def generate_text():
 
     # Decode the generated text
     decoded_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    # Return generated text as JSON response
-    return jsonify({"generated_text": decoded_text})
-
-# Run the Flask app
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000)
+    return decoded_text
+
+# Set up Gradio interface
+interface = gr.Interface(
+    fn=generate_text,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your text here..."),
+    outputs="text",
+    title="Text Generator",
+    description="Generate text using the Lumenspark model."
+)
+
+# Launch the interface
+interface.launch()