Update launch.py

launch.py

@@ -1,106 +1,64 @@
 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
+def main():
+    # Use the current directory where the model files are located
+    model_path = os.path.dirname(os.path.abspath(__file__))
+    
+    print("Loading model and tokenizer...")
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        
+        # Load the model. trust_remote_code=True is included in case of custom architectures
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            trust_remote_code=True,
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+        )
+        
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        model.to(device)
+        print(f"Model loaded successfully on {device.upper()}!\n")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return
 
-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=256, 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...")
+    print("--- Chat Session Started (Type 'exit' or 'quit' to end) ---")
+    
+    while True:
+        try:
+            user_input = input("You: ").strip()
+            if not user_input:
+                continue
+            if user_input.lower() in ["exit", "quit"]:
+                print("Ending session. Goodbye!")
                 break
+            
+            # Tokenize input
+            inputs = tokenizer(user_input, return_tensors="pt").to(device)
+            
+            # Generate response
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=128,
+                    do_sample=True,
+                    temperature=0.6,
+                    top_k=40,
+                    top_p=0.9,
+                    pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+                )
+            
+            # Decode only the newly generated tokens
+            input_length = inputs.input_ids.shape[1]
+            response = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+            
+            print(f"Model: {response.strip()}\n")
+            
+        except KeyboardInterrupt:
+            print("\nSession interrupted. Goodbye!")
+            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.")
+    main()