Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from theme import fast_rtc_theme | |
| import torch | |
| import json | |
| import uuid | |
| import os | |
| import time | |
| import pytz | |
| from datetime import datetime | |
| from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer | |
| from threading import Thread | |
| from huggingface_hub import CommitScheduler | |
| from pathlib import Path | |
| import spaces | |
| os.system("apt-get update && apt-get install -y libstdc++6") | |
| token = os.environ["HF_TOKEN"] | |
| model_id = "large-traversaal/Mantra-14B" | |
| model = AutoModelForCausalLM.from_pretrained(model_id, token=token, trust_remote_code=True, torch_dtype=torch.bfloat16) | |
| tokenizer = AutoTokenizer.from_pretrained(model_id, token=token) | |
| terminators = [tokenizer.eos_token_id] | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| model = model.to(device) | |
| log_folder = Path("logs") | |
| log_folder.mkdir(parents=True, exist_ok=True) | |
| log_file = log_folder / f"chat_log_{uuid.uuid4()}.json" | |
| scheduler = CommitScheduler(repo_id="large-traversaal/mantra-14b-user-interaction-log", repo_type="dataset", folder_path=log_folder, path_in_repo="data", every=0.01, token=token) | |
| timezone = pytz.timezone("UTC") | |
| def chat(message, history, temperature, do_sample, max_tokens, top_p): | |
| start_time = time.time() | |
| timestamp = datetime.now(timezone).strftime("%Y-%m-%d %H:%M:%S %Z") | |
| conversation_history = [] | |
| for item in history: | |
| conversation_history.append({"role": "user", "content": item[0]}) | |
| if item[1] is not None: | |
| conversation_history.append({"role": "assistant", "content": item[1]}) | |
| conversation_history.append({"role": "user", "content": message}) | |
| messages = tokenizer.apply_chat_template(conversation_history, tokenize=False, add_generation_prompt=True) | |
| model_inputs = tokenizer([messages], return_tensors="pt").to(device) | |
| streamer = TextIteratorStreamer(tokenizer, timeout=70.0, skip_prompt=True, skip_special_tokens=True) | |
| generate_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=max_tokens, do_sample=do_sample, temperature=temperature, top_p=top_p, eos_token_id=terminators,) | |
| if temperature == 0: | |
| generate_kwargs["do_sample"] = False | |
| generation_thread = Thread(target=model.generate, kwargs=generate_kwargs) | |
| generation_thread.start() | |
| partial_text = "" | |
| for new_text in streamer: | |
| partial_text += new_text | |
| yield partial_text | |
| response_time = round(time.time() - start_time, 2) | |
| log_data = {"timestamp": timestamp,"input": message,"output": partial_text,"response_time": response_time,"temperature": temperature,"do_sample": do_sample,"max_tokens": max_tokens,"top_p": top_p} | |
| with scheduler.lock: | |
| with log_file.open("a", encoding="utf-8") as f: | |
| f.write(json.dumps(log_data, ensure_ascii=False) + "\n") | |
| def clear_chat(): | |
| return [], [] | |
| def export_chat(history): | |
| if not history: | |
| return None # No chat history to export | |
| file_path = "chat_history.txt" | |
| with open(file_path, "w", encoding="utf-8") as f: | |
| for msg in history: | |
| f.write(f"User: {msg[0]}\nBot: {msg[1]}\n") | |
| return file_path | |
| with gr.Blocks(theme=fast_rtc_theme) as demo: | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("#### ⚙️🛠 Configure Settings") | |
| temperature = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.1, label="Temperature", interactive=True) | |
| do_sample = gr.Checkbox(label="Sampling", value=True, interactive=True) | |
| max_tokens = gr.Slider(minimum=128, maximum=4096, step=1, value=1024, label="max_new_tokens", interactive=True) | |
| top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.1, step=0.2, label="top_p", interactive=True) | |
| with gr.Column(scale=3): | |
| gr.Markdown("# **Chat With Mantra-14B** 💬 ") | |
| chat_interface = gr.ChatInterface(fn=chat, | |
| examples=[["What is the English translation of: 'इस मॉडल को हिंदी और अंग्रेजी डेटा पर प्रशिक्षित किया गया था'?"], | |
| ["टिम अपने 3 बच्चों को ट्रिक या ट्रीटिंग के लिए ले जाता है। वे 4 घंटे बाहर रहते हैं। हर घंटे वे x घरों में जाते हैं। हर घर में हर बच्चे को 3 ट्रीट मिलते हैं। उसके बच्चों को कुल 180 ट्रीट मिलते हैं। अज्ञात चर x का मान क्या है?"], | |
| ["How do you play fetch? A) Throw the object for the dog to bring back to you. B) Get the object and bring it back to the dog."]], | |
| additional_inputs=[temperature, do_sample, max_tokens, top_p], | |
| stop_btn="⏹ Stop", | |
| description="Mantra-14B is a bilingual instruction-tuned LLM for Hindi and English, trained on a mixed datasets composed of 485K Hindi-English samples.",) | |
| with gr.Row(): | |
| clear_btn = gr.Button("🧹 Clear Chat", variant="primary") | |
| export_btn = gr.Button("📥 Export Chat", variant="primary") | |
| clear_btn.click(fn=clear_chat, outputs=[chat_interface.chatbot, chat_interface.chatbot_value]) | |
| export_btn.click(fn=export_chat, inputs=[chat_interface.chatbot], outputs=[gr.File()]) | |
| demo.launch() |