Spaces:
Sleeping
Sleeping
File size: 1,562 Bytes
647051c 76ecd6c 647051c 4956caf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import os
import shutil
from flask import Flask, render_template, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
app = Flask(__name__)
# Kill any process using port 7860
def kill_port(port):
for proc in psutil.process_iter(attrs=["pid", "connections"]):
for conn in proc.info["connections"]:
if conn.laddr.port == port:
os.kill(proc.info["pid"], 9)
kill_port(7860) # Ensure Flask doesn't crash due to a used port
# Define cache directory
os.environ["HF_HOME"] = "/app/cache"
# Load Myanmarsar-GPT (1.42B params) from Hugging Face
MODEL_NAME = "simbolo-ai/Myanmarsar-GPT"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=cache_dir)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, cache_dir=cache_dir)
# Function to generate chatbot responses
def generate_response(prompt):
inputs = tokenizer(prompt, return_tensors="pt")
with torch.no_grad():
output = model.generate(**inputs, max_length=200)
return tokenizer.decode(output[0], skip_special_tokens=True)
# Serve the HTML page
@app.route("/")
def home():
return render_template("index.html")
# API route for chatbot responses
@app.route("/chat", methods=["POST"])
def chat():
user_message = request.json.get("message", "")
bot_reply = generate_response(user_message)
return jsonify({"reply": bot_reply})
if __name__ == "__main__":
port = int(os.environ.get("PORT", 7860)) # Default to 7860, but use any assigned port
app.run(host="0.0.0.0", port=port) |