Smllm / app2.py
ghosthets's picture
Rename app.py to app2.py
b50fca9 verified
from flask import Flask, request, jsonify
from transformers import pipeline
import torch
app = Flask(__name__)
# ===========================
# LOAD LOCAL/FREE MODEL (HF TOKEN NAHI CHAHIYE)
# ===========================
# Apna model ID yaha daal sakte ho, ya HuggingFace free model
model_id = "google/gemma-2b" # Example free model
print("🔄 Loading model...")
# CPU/GPU device set
device = 0 if torch.cuda.is_available() else -1
# pipeline me device=-1 -> CPU, device=0 -> GPU
ai = pipeline("text-generation", model=model_id, max_new_tokens=200, device=device)
print("✅ Model loaded!")
# ===========================
# CHAT API
# ===========================
@app.route('/chat', methods=['POST'])
def chat():
try:
data = request.get_json()
msg = data.get("message", "")
if not msg:
return jsonify({"error": "No message sent"}), 400
output = ai(msg)[0]["generated_text"]
return jsonify({"reply": output})
except Exception as e:
return jsonify({"error": str(e)}), 500
# ===========================
# RUN SERVER
# ===========================
if __name__ == "__main__":
app.run(host='0.0.0.0', port=7860)