Spaces:
Running
Running
File size: 3,297 Bytes
1d5b573 928cd07 1d5b573 928cd07 43f1789 928cd07 43f1789 928cd07 43f1789 928cd07 1d5b573 928cd07 1d5b573 928cd07 1d5b573 |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
from flask import Flask, request, jsonify
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import os
app = Flask(__name__)
# Configuraci贸n CORS para permitir solicitudes desde tu dominio
@app.after_request
def after_request(response):
response.headers.add('Access-Control-Allow-Origin', 'https://justbyte.es')
response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
response.headers.add('Access-Control-Allow-Methods', 'POST')
return response
# Cargar el modelo y tokenizador (se carga una sola vez al iniciar)
@app.before_first_request
def load_model():
global model, tokenizer
print("Cargando modelo Zephyr-7B...")
# Cargar el modelo y tokenizador
model_name = "HuggingFaceH4/zephyr-7b-beta"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16, # Usar precisi贸n reducida para ahorrar memoria
device_map="auto", # Distribuir el modelo autom谩ticamente
load_in_8bit=True # Cuantizaci贸n a 8 bits para reducir uso de memoria
)
print("Modelo cargado correctamente!")
# Cargar el prompt desde el archivo
def get_system_prompt():
with open("prompt.txt", "r", encoding="utf-8") as f:
return f.read().strip()
@app.route('/generate', methods=['POST'])
def generate_response():
try:
# Obtener la pregunta del usuario
data = request.json
user_message = data.get('message', '')
if not user_message:
return jsonify({"error": "No se proporcion贸 ninguna pregunta"}), 400
# Obtener el prompt del sistema
system_prompt = get_system_prompt()
# Crear el formato de conversaci贸n para Zephyr
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_message}
]
# Convertir mensajes al formato que espera el modelo
prompt = tokenizer.apply_chat_template(messages, tokenize=False)
# Generar respuesta
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
# Configuraci贸n de generaci贸n
generation_config = {
"max_new_tokens": 500,
"temperature": 0.7,
"top_p": 0.9,
"do_sample": True,
"pad_token_id": tokenizer.eos_token_id
}
# Generar respuesta
with torch.no_grad():
outputs = model.generate(**inputs, **generation_config)
# Decodificar la respuesta
full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
# Extraer solo la respuesta del asistente (despu茅s del 煤ltimo mensaje del usuario)
assistant_response = full_response.split("assistant:")[-1].strip()
return jsonify({"response": assistant_response})
except Exception as e:
print(f"Error: {str(e)}")
return jsonify({"error": f"Error al generar respuesta: {str(e)}"}), 500
if __name__ == '__main__':
# Puerto que Hugging Face Spaces utiliza
port = int(os.environ.get('PORT', 7860))
app.run(host='0.0.0.0', port=port)
|