SINTEIX / RDX-7b-V1.py
lennych's picture
Upload RDX-7b-V1.py
57e758c
import transformers
model_name = 'Intel/neural-chat-7b-v3-1'
try:
model = transformers.AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
except Exception as e:
print(f"Error al cargar el modelo: {e}")
def generate_response(system_input, user_input, max_length=1000, num_return_sequences=1):
"""
Genera una respuesta del asistente dada una entrada del sistema y del usuario.
Parámetros:
system_input (str): La entrada del sistema.
user_input (str): La entrada del usuario.
max_length (int, opcional): La longitud máxima de la respuesta generada. Por defecto es 1000.
num_return_sequences (int, opcional): El número de secuencias de retorno. Por defecto es 1.
Devuelve:
str: La respuesta generada del asistente.
"""
# Formatear la entrada usando la plantilla proporcionada
prompt = f"### System:\n{system_input}\n### User:\n{user_input}\n### Assistant:\n"
# Tokenizar y codificar el prompt
inputs = tokenizer.encode(prompt, return_tensors="pt", add_special_tokens=False)
# Generar una respuesta
try:
outputs = model.generate(inputs, max_length=max_length, num_return_sequences=num_return_sequences)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
except Exception as e:
print(f"Error al generar la respuesta: {e}")
return None
# Extraer solo la respuesta del asistente
return response.split("### Assistant:\n")[-1]