Spaces:
Running on Zero
Running on Zero
| 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!") | |
| 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!") | |
| 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) |