Neura-Tech-AI's picture
Update app.py
4003c14 verified
Raw
History Blame Contribute Delete
4.48 kB
import json
from datetime import datetime
import os
from huggingface_hub import HfApi
import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Neura-Tech-AI/Neuron-V1-14B-Instruct"
LOG_FILE = "chat_logs.jsonl"
tokenizer = None
model = None
def save_chat(user_msg, assistant_msg):
print("Saving chat...")
log = {
"timestamp": datetime.utcnow().isoformat(),
"user": user_msg,
"assistant": assistant_msg,
}
with open(LOG_FILE, "a", encoding="utf-8") as f:
f.write(json.dumps(log, ensure_ascii=False) + "\n")
print("Chat saved!")
@spaces.GPU
def load_model():
global tokenizer, model
if model is not None:
return
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
MODEL_ID,
trust_remote_code=True,
)
print("Loading model...")
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True,
trust_remote_code=True,
)
model.eval()
print("Model loaded successfully!")
@spaces.GPU
def chat(
message,
history,
system_prompt,
max_new_tokens,
temperature,
top_p,
):
load_model()
messages = []
if system_prompt.strip():
messages.append(
{
"role": "system",
"content": system_prompt,
}
)
if history:
for user_msg, assistant_msg in history:
messages.append(
{
"role": "user",
"content": user_msg,
}
)
messages.append(
{
"role": "assistant",
"content": assistant_msg,
}
)
messages.append(
{
"role": "user",
"content": message,
}
)
prompt = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
inputs = tokenizer(
prompt,
return_tensors="pt",
)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
with torch.inference_mode():
outputs = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
)
generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]
response = tokenizer.decode(
generated_tokens,
skip_special_tokens=True,
).strip()
# Save chat log
save_chat(message, response)
return response
demo = gr.ChatInterface(
fn=chat,
title="Neuron V1 14B Instruct",
description="Developed by Neura Tech AI",
additional_inputs=[
gr.Textbox(
value="You are Neuron, developed by Neura Tech AI.",
label="System Prompt",
),
gr.Slider(
minimum=1,
maximum=2048,
value=512,
step=1,
label="Max New Tokens",
),
gr.Slider(
minimum=0.0,
maximum=2.0,
value=0.7,
step=0.1,
label="Temperature",
),
gr.Slider(
minimum=0.1,
maximum=1.0,
value=0.95,
step=0.05,
label="Top-p",
),
],
)
demo.queue(max_size=32)
if __name__ == "__main__":
demo.launch()
import os
LOG_FILE = "chat_logs.jsonl"
DATASET_REPO = "Neura-Tech-AI/chat-logs"
def save_chat(user_msg, assistant_msg):
log = {
"timestamp": datetime.utcnow().isoformat(),
"user": user_msg,
"assistant": assistant_msg,
}
# Local file me append
with open(LOG_FILE, "a", encoding="utf-8") as f:
f.write(json.dumps(log, ensure_ascii=False) + "\n")
# HF Dataset par upload
token = os.getenv("HF_TOKEN")
if token:
try:
api = HfApi(token=token)
api.upload_file(
path_or_fileobj=LOG_FILE,
path_in_repo="chat_logs.jsonl",
repo_id=DATASET_REPO,
repo_type="dataset",
)
print("Uploaded logs to dataset.")
except Exception as e:
print("Upload failed:", e)