Update app.py

app.py

@@ -1,9 +1,15 @@
 from typing import List, Tuple, Dict, Generator
-from langchain.llms import OpenAI
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import gradio as gr
 
-model_name = "gpt-3.5-turbo"
-LLM = OpenAI(model_name=model_name, temperature=0.1)
+# Load your safetensors model and tokenizer
+model_name = "DuckyPolice/ElapticAI-1a"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
 
 def create_history_messages(history: List[Tuple[str, str]]) -> List[dict]:
     history_messages = [{"role": "user", "content": m[0]} for m in history]
@@ -36,9 +42,7 @@ def create_formatted_history(history_messages: List[dict]) -> List[Tuple[str, st
 
     return formatted_history
 
-def chat(
-    message: str, state: List[Dict[str, str]], client = LLM.client
-) -> Generator[Tuple[List[Tuple[str, str]], List[Dict[str, str]]], None, None]:
+def chat(message: str, state: List[Dict[str, str]]) -> Generator[Tuple[List[Tuple[str, str]], List[Dict[str, str]]], None, None]:
     history_messages = state
     if history_messages == None:
         history_messages = []
@@ -48,36 +52,25 @@ def chat(
     # We have no content for the assistant's response yet but we will update this:
     history_messages.append({"role": "assistant", "content": ""})
 
-    response_message = ""
+    # Prepare input for the model
+    inputs = tokenizer(message, return_tensors="pt").to(device)
 
-    chat_generator = client.create(
-        messages=history_messages, stream=True, model=model_name
-    )
+    # Generate response from model
+    response_ids = model.generate(inputs['input_ids'], max_length=200)
+    response_message = tokenizer.decode(response_ids[0], skip_special_tokens=True)
 
-    for chunk in chat_generator:
-        if "choices" in chunk:
-            for choice in chunk["choices"]:
-                if "delta" in choice and "content" in choice["delta"]:
-                    new_token = choice["delta"]["content"]
-                    # Add the latest token:
-                    response_message += new_token
-                    # Update the assistant's response in our model:
-                    history_messages[-1]["content"] = response_message
+    # Update the assistant's response in our model
+    history_messages[-1]["content"] = response_message
 
-                if "finish_reason" in choice and choice["finish_reason"] == "stop":
-                    break
-        formatted_history = create_formatted_history(history_messages)
-        yield formatted_history, history_messages
+    formatted_history = create_formatted_history(history_messages)
+    yield formatted_history, history_messages
 
 chatbot = gr.Chatbot(label="Chat").style(color_map=("yellow", "purple"))
 iface = gr.Interface(
     fn=chat,
-    inputs=[
-        gr.Textbox(placeholder="Hello! How are you? etc.", label="Message"),
-        "state",
-    ],
+    inputs=[gr.Textbox(placeholder="Hello! How are you? etc.", label="Message"), "state"],
     outputs=[chatbot, "state"],
     allow_flagging="never",
 )
 
-iface.queue().launch()
+iface.queue().launch()