from flask import Flask, request, jsonify, send_from_directory
import torch
from transformers import pipeline

# Khởi tạo ứng dụng Flask
app = Flask(__name__)

# Khởi tạo mô hình
model_id = "meta-llama/Llama-3.2-3B-Instruct"
pipe = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16)

@app.route('/')
def home():
    return send_from_directory('.', 'index.html')

@app.route('/style.css')
def style():
    return send_from_directory('.', 'style.css')

@app.route('/chat', methods=['POST'])
def chat():
    user_input = request.json.get('message')
    response = pipe(user_input, max_new_tokens=256)
    return jsonify({'response': response[0]['generated_text']})

if __name__ == '__main__':
    app.run()