askbyte commited on
Commit
6996ba5
verified
1 Parent(s): 0fa6e0f

Create app_gemma.py

Browse files
Files changed (1) hide show
  1. app_gemma.py +99 -0
app_gemma.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Versi贸n usando Gemma-2B, un modelo peque帽o de Google.
3
+ Renombra este archivo a app.py si prefieres usar este modelo.
4
+ """
5
+
6
+ import gradio as gr
7
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
8
+ import torch
9
+ import gc
10
+ import os
11
+
12
+ # Funci贸n para cargar el prompt desde el archivo
13
+ def get_system_prompt():
14
+ with open("prompt.txt", "r", encoding="utf-8") as f:
15
+ return f.read().strip()
16
+
17
+ # Funci贸n para generar respuestas
18
+ def generate_response(user_message):
19
+ try:
20
+ if not user_message.strip():
21
+ return "Por favor, escribe una pregunta para que pueda ayudarte."
22
+
23
+ # Obtener el prompt del sistema
24
+ system_prompt = get_system_prompt()
25
+
26
+ # Crear el prompt completo para Gemma
27
+ 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"
28
+
29
+ # Generar respuesta usando el pipeline
30
+ response = generator(
31
+ prompt,
32
+ max_new_tokens=256, # Reducido para mayor velocidad
33
+ temperature=0.7,
34
+ top_p=0.9,
35
+ do_sample=True,
36
+ num_return_sequences=1
37
+ )[0]["generated_text"]
38
+
39
+ # Extraer solo la respuesta del modelo (despu茅s del 煤ltimo <start_of_turn>model)
40
+ assistant_response = response.split("<start_of_turn>model\n")[-1].split("<end_of_turn>")[0].strip()
41
+
42
+ # Forzar recolecci贸n de basura
43
+ gc.collect()
44
+ torch.cuda.empty_cache() if torch.cuda.is_available() else None
45
+
46
+ return assistant_response
47
+
48
+ except Exception as e:
49
+ print(f"Error: {str(e)}")
50
+ return f"Lo siento, ha ocurrido un error: {str(e)}"
51
+
52
+ # Cargar el modelo y crear el pipeline
53
+ print("Cargando modelo Gemma-2B...")
54
+
55
+ # Usar Gemma-2B, un modelo peque帽o (2B par谩metros) pero potente
56
+ model_name = "google/gemma-2b-it"
57
+
58
+ # Configuraci贸n para reducir el uso de memoria y aumentar velocidad
59
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
60
+ model = AutoModelForCausalLM.from_pretrained(
61
+ model_name,
62
+ torch_dtype=torch.float32,
63
+ low_cpu_mem_usage=True,
64
+ # Cargamos en 8-bit para reducir memoria y aumentar velocidad
65
+ load_in_8bit=True
66
+ )
67
+
68
+ # Crear el pipeline de generaci贸n de texto
69
+ generator = pipeline(
70
+ "text-generation",
71
+ model=model,
72
+ tokenizer=tokenizer
73
+ )
74
+
75
+ print("Modelo cargado correctamente!")
76
+
77
+ # Crear la interfaz de Gradio
78
+ demo = gr.Interface(
79
+ fn=generate_response,
80
+ inputs=gr.Textbox(
81
+ placeholder="Escribe tu pregunta sobre emprendimiento aqu铆...",
82
+ label="Tu pregunta"
83
+ ),
84
+ outputs=gr.Textbox(label="Respuesta de BITER"),
85
+ title="BITER - Tu Mentor en Tiempo Real para Decisiones de Negocio",
86
+ description="BITER es un asistente de IA que responde dudas de emprendedores como si fuera un CEO experimentado.",
87
+ examples=[
88
+ ["驴C贸mo puedo validar mi idea de negocio con poco presupuesto?"],
89
+ ["驴Cu谩l es la mejor estrategia para conseguir mis primeros clientes?"],
90
+ ["驴Deber铆a invertir en publicidad en redes sociales o en SEO?"]
91
+ ],
92
+ allow_flagging="never"
93
+ )
94
+
95
+ # Lanzar la aplicaci贸n con configuraci贸n para ahorrar memoria
96
+ if __name__ == "__main__":
97
+ # Configurar menos workers para ahorrar memoria y aumentar velocidad
98
+ demo.queue(max_size=1).launch(share=False, debug=False)
99
+