Spaces:
Running
Running
""" | |
Versi贸n usando Gemma-2B, un modelo peque帽o de Google. | |
Renombra este archivo a app.py si prefieres usar este modelo. | |
""" | |
import gradio as gr | |
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline | |
import torch | |
import gc | |
import os | |
# Funci贸n para cargar el prompt desde el archivo | |
def get_system_prompt(): | |
with open("prompt.txt", "r", encoding="utf-8") as f: | |
return f.read().strip() | |
# Funci贸n para generar respuestas | |
def generate_response(user_message): | |
try: | |
if not user_message.strip(): | |
return "Por favor, escribe una pregunta para que pueda ayudarte." | |
# Obtener el prompt del sistema | |
system_prompt = get_system_prompt() | |
# Crear el prompt completo para Gemma | |
prompt = f"<start_of_turn>system\n{system_prompt}<end_of_turn>\n<start_of_turn>user\n{user_message}<end_of_turn>\n<start_of_turn>model\n" | |
# Generar respuesta usando el pipeline | |
response = generator( | |
prompt, | |
max_new_tokens=256, # Reducido para mayor velocidad | |
temperature=0.7, | |
top_p=0.9, | |
do_sample=True, | |
num_return_sequences=1 | |
)[0]["generated_text"] | |
# Extraer solo la respuesta del modelo (despu茅s del 煤ltimo <start_of_turn>model) | |
assistant_response = response.split("<start_of_turn>model\n")[-1].split("<end_of_turn>")[0].strip() | |
# Forzar recolecci贸n de basura | |
gc.collect() | |
torch.cuda.empty_cache() if torch.cuda.is_available() else None | |
return assistant_response | |
except Exception as e: | |
print(f"Error: {str(e)}") | |
return f"Lo siento, ha ocurrido un error: {str(e)}" | |
# Cargar el modelo y crear el pipeline | |
print("Cargando modelo Gemma-2B...") | |
# Usar Gemma-2B, un modelo peque帽o (2B par谩metros) pero potente | |
model_name = "google/gemma-2b-it" | |
# Configuraci贸n para reducir el uso de memoria y aumentar velocidad | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_name, | |
torch_dtype=torch.float32, | |
low_cpu_mem_usage=True, | |
# Cargamos en 8-bit para reducir memoria y aumentar velocidad | |
load_in_8bit=True | |
) | |
# Crear el pipeline de generaci贸n de texto | |
generator = pipeline( | |
"text-generation", | |
model=model, | |
tokenizer=tokenizer | |
) | |
print("Modelo cargado correctamente!") | |
# Crear la interfaz de Gradio | |
demo = gr.Interface( | |
fn=generate_response, | |
inputs=gr.Textbox( | |
placeholder="Escribe tu pregunta sobre emprendimiento aqu铆...", | |
label="Tu pregunta" | |
), | |
outputs=gr.Textbox(label="Respuesta de BITER"), | |
title="BITER - Tu Mentor en Tiempo Real para Decisiones de Negocio", | |
description="BITER es un asistente de IA que responde dudas de emprendedores como si fuera un CEO experimentado.", | |
examples=[ | |
["驴C贸mo puedo validar mi idea de negocio con poco presupuesto?"], | |
["驴Cu谩l es la mejor estrategia para conseguir mis primeros clientes?"], | |
["驴Deber铆a invertir en publicidad en redes sociales o en SEO?"] | |
], | |
allow_flagging="never" | |
) | |
# Lanzar la aplicaci贸n con configuraci贸n para ahorrar memoria | |
if __name__ == "__main__": | |
# Configurar menos workers para ahorrar memoria y aumentar velocidad | |
demo.queue(max_size=1).launch(share=False, debug=False) | |