Spaces:

helloperson123
/

tiny-llama-chatbot

Runtime error

helloperson123 commited on Oct 10

Commit

e08d99f

verified ·

1 Parent(s): 437a17e

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -1,44 +1,28 @@
-from transformers import pipeline
-from flask import Flask, request, render_template_string
 app = Flask(__name__)
-pipe = pipeline("text-generation", model="sshleifer/tiny-gpt2")
-HTML_PAGE = """
-<!DOCTYPE html>
-<html>
-  <head>
-    <title>TinyGPT2 Chat</title>
-    <style>
-      body { font-family: sans-serif; margin: 40px; }
-      textarea { width: 100%; height: 100px; }
-      button { margin-top: 10px; padding: 8px 16px; }
-      .output { margin-top: 20px; white-space: pre-wrap; }
-    </style>
-  </head>
-  <body>
-    <h1>🤖 TinyGPT2 Chat</h1>
-    <form method="POST">
-      <textarea name="prompt" placeholder="Type your message here...">{{prompt}}</textarea><br>
-      <button type="submit">Generate</button>
-    </form>
-    {% if output %}
-      <div class="output"><strong>AI:</strong> {{output}}</div>
-    {% endif %}
-  </body>
-</html>
-"""
-@app.route("/", methods=["GET", "POST"])
-def chat():
-    output = ""
-    prompt = ""
-    if request.method == "POST":
-        prompt = request.form["prompt"]
-        result = pipe(prompt, max_length=100)[0]["generated_text"]
-        output = result
-    return render_template_string(HTML_PAGE, output=output, prompt=prompt)
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)

+from flask import Flask, request, jsonify
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 app = Flask(__name__)
+# Load model and tokenizer
+MODEL_NAME = "openai-community/gpt2"  # or your Tiny LLaMA model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
+@app.route("/api/ask", methods=["POST"])
+def ask():
+    data = request.get_json()
+    prompt = data.get("prompt", "")
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=50)
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return jsonify({"reply": text})
+@app.route("/")
+def home():
+    return "✅ Model API running! POST JSON to /api/ask with {'prompt': 'your text'}"
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)