Update app.py

app.py

@@ -1,31 +1,51 @@
+import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Load the HelpingAI2.5-2B model
 model = AutoModelForCausalLM.from_pretrained("OEvortex/HelpingAI2.5-2B")
-# Load the tokenizer
 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,
+# 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 structure
+    chat = [
+        { "role": "system", "content": "You are HelpingAI, an emotional AI. Always answer my questions in the HelpingAI style." },
+        { "role": "user", "content": user_input }
+    ]
+
+    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,
+    )
+
+    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,
+    inputs="text",
+    outputs="text",
+    live=True
 )
 
-response = outputs[0][inputs.shape[-1]:]
-print(tokenizer.decode(response, skip_special_tokens=True))
+# Launch the Gradio app
+iface.launch()