Update app.py

app.py

@@ -1,54 +1,31 @@
-import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the HelpingAI2.5-10B model
-model = AutoModelForCausalLM.from_pretrained("HelpingAI/HelpingAI-flash")
+# Load the HelpingAI2.5-2B model
+model = AutoModelForCausalLM.from_pretrained("OEvortex/HelpingAI2.5-2B")
 # Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("HelpingAI/HelpingAI-flash")
-
-# Move model to GPU (if available) or CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-# Define the function for generating responses
-def generate_response(user_input):
-    # Define the chat input
-    chat = [
-        { "role": "system", "content": "You are HelpingAI, an emotional AI. Always answer my questions in the HelpingAI style." },
-        { "role": "user", "content": user_input }
-    ]
-
-    # Concatenate the chat conversation into a single string with the proper format
-    chat_input = ""
-    for message in chat:
-        role = message["role"]
-        content = message["content"]
-        chat_input += f"{role}: {content}\n"
-
-    # Tokenize the input
-    inputs = tokenizer(chat_input, return_tensors="pt").to(device)
-
-    # Generate text
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=256,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9,
-    )
-
-    # Decode the generated text
-    response = outputs[0][inputs["input_ids"].shape[-1]:]
-    return tokenizer.decode(response, skip_special_tokens=True)
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=generate_response,  # Function to generate response
-    inputs="text",          # Input is a text box for the user to type in
-    outputs="text",         # Output is the generated response (text)
-    live=True               # Allows updating as the user types
+tokenizer = AutoTokenizer.from_pretrained("OEvortex/HelpingAI2.5-2B")
+
+# Define the chat input
+chat = [
+    { "role": "system", "content": "You are HelpingAI, an emotional AI. Always answer my questions in the HelpingAI style." },
+    { "role": "user", "content": "GIVE ME YOUR INTRO" }
+]
+
+inputs = tokenizer.apply_chat_template(
+    chat,
+    add_generation_prompt=True,
+    return_tensors="pt"
+).to(model.device)
+
+# Generate text
+outputs = model.generate(
+    inputs,
+    max_new_tokens=256,
+    do_sample=True,
+    temperature=0.6,
+    top_p=0.9,
 )
 
-# Launch the Gradio app
-iface.launch()
+response = outputs[0][inputs.shape[-1]:]
+print(tokenizer.decode(response, skip_special_tokens=True))