| import os
|
| import sys
|
| import torch
|
| from transformers import AutoModelForCausalLM, AutoTokenizer
|
| import glob
|
| from prompt_toolkit import PromptSession
|
| from prompt_toolkit.keys import Keys
|
| from prompt_toolkit.application import Application
|
| from prompt_toolkit.layout import Layout, HSplit, Window
|
| from prompt_toolkit.widgets import Frame, Label
|
| from prompt_toolkit.key_binding import KeyBindings
|
| from prompt_toolkit.styles import Style
|
|
|
| class DummyTorchAudio:
|
| pass
|
|
|
| sys.modules['torchaudio'] = DummyTorchAudio()
|
|
|
| def get_model_files():
|
| return glob.glob("*.safetensors")
|
|
|
| def select_file_with_prompt(files):
|
| if not files:
|
| print("No .safetensors files found.")
|
| return None
|
| if len(files) == 1:
|
| return files[0]
|
| selected = 0
|
| style = Style.from_dict({"selected": "fg:ansiblack bg:ansigreen"})
|
| kb = KeyBindings()
|
|
|
| @kb.add("up")
|
| def up(event):
|
| nonlocal selected
|
| selected = (selected - 1) % len(files)
|
|
|
| @kb.add("down")
|
| def down(event):
|
| nonlocal selected
|
| selected = (selected + 1) % len(files)
|
|
|
| @kb.add("enter")
|
| def enter(event):
|
| event.app.exit(result=files[selected])
|
|
|
| @kb.add("c-c")
|
| def cancel(event):
|
| event.app.exit(result=None)
|
|
|
| body = HSplit([])
|
| for idx, file in enumerate(files):
|
| if idx == selected:
|
| body.children.append(Window(Label(text=f"[x] {file}"), style="selected"))
|
| else:
|
| body.children.append(Window(Label(text=f"[ ] {file}")))
|
|
|
| container = Frame(body=body, title="Select model (up/down, enter)")
|
| layout = Layout(container)
|
| app = Application(layout=layout, key_bindings=kb, style=style, full_screen=False)
|
| return app.run()
|
|
|
| class ChatClient:
|
| def __init__(self, model_path):
|
| self.tokenizer = AutoTokenizer.from_pretrained(".", local_files_only=True)
|
| self.model = AutoModelForCausalLM.from_pretrained(".", local_files_only=True)
|
| self.history = []
|
|
|
| def generate_response(self, prompt):
|
| inputs = self.tokenizer(prompt, return_tensors="pt")
|
| with torch.no_grad():
|
| outputs = self.model.generate(**inputs, max_new_tokens=128, do_sample=True, temperature=0.7)
|
| response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
|
| response = response[len(prompt):].strip()
|
| return response
|
|
|
| def chat(self):
|
| session = PromptSession()
|
| print("Chat with model. Type 'exit' to quit.")
|
| print("Type 'clear' to clear history.")
|
| while True:
|
| try:
|
| user_input = session.prompt("You: ")
|
| if user_input.lower() in ["exit", "quit"]:
|
| break
|
| if user_input.lower() == "clear":
|
| self.history = []
|
| print("History cleared.")
|
| continue
|
| self.history.append("User: " + user_input)
|
| context = "\n".join(self.history[-10:])
|
| response = self.generate_response(context + "\nAssistant:")
|
| print("Assistant:", response)
|
| self.history.append("Assistant: " + response)
|
| except KeyboardInterrupt:
|
| print("\nExiting...")
|
| break
|
|
|
| if __name__ == "__main__":
|
| files = get_model_files()
|
| selected_file = select_file_with_prompt(files)
|
| if selected_file:
|
| print(f"Loading {selected_file}...")
|
| client = ChatClient(selected_file)
|
| client.chat()
|
| else:
|
| print("No model selected.") |
