Update app.py

app.py

@@ -2,69 +2,75 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Set the device to CPU since Hugging Face Spaces does not support GPU
-device = torch.device("cpu")  # Ensure it's using CPU only
+# Define device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Load model and tokenizer
-model_name = "hosseinhimself/ISANG-v1.0-8B"  # Replace with your model name
+# Load the model and tokenizer
+model_name = "hosseinhimself/ISANG-v1.0-8B"
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
 model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
-# Define the Alpaca-style prompt template
-alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
-### Instruction:
-You are ISANG, a multilingual large language model made by ISANG AI. You only respond in Persian, Korean, or English. If a user uses one of these languages, reply in the same language.
-
-### Input:
-{}
-
-### Response:
-{}"""
-
-# Function to generate responses
-def generate_response(input_text, max_tokens=1024, temperature=0.7, history=[]):
-    # Retain only the last two exchanges for context
-    if len(history) > 2:
-        history = history[-2:]
-    
-    # Format the prompt
-    prompt = "\n".join(history + [f"User: {input_text}\nAI:"])
-    
-    # Tokenize the input
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    
-    # Generate model output
-    output = model.generate(
-        inputs.input_ids,
-        max_new_tokens=max_tokens,
-        temperature=temperature
-    )
-    
-    # Decode the model output
-    response = tokenizer.decode(output[0], skip_special_tokens=True).strip()
-    
-    # Update the history
-    history.append(f"User: {input_text}")
-    history.append(f"AI: {response}")
-    
-    return response, history
-
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_response,
-    inputs=[
-        gr.Textbox(label="Your Message", placeholder="Type your message here..."),
-        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max Tokens"),
-        gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
-        gr.State(value=[])  # To maintain conversation history
-    ],
-    outputs=[gr.Textbox(label="AI Response"), gr.State()],
-    title="ISANG Chatbot",
-    description="A chatbot powered by ISANG-v1.0-8B model. Chat with me!",
-    theme="huggingface",  # Purple theme
-    live=False  # Set to False since live updates aren't required
+def chat_with_model(history, user_input):
+    """
+    Generate a response using the model, considering the last two interactions.
+
+    Parameters:
+        history (list of tuples): Conversation history as a list of (user, bot) pairs.
+        user_input (str): The latest user input.
+
+    Returns:
+        history (list of tuples): Updated conversation history.
+    """
+    # Use the last two interactions for context
+    context = ""
+    for user_message, bot_message in history[-2:]:
+        context += f"User: {user_message}\nBot: {bot_message}\n"
+
+    # Add the current user input
+    context += f"User: {user_input}\nBot:"
+
+    # Tokenize and generate a response
+    inputs = tokenizer(context, return_tensors="pt", truncation=True).to(device)
+    output = model.generate(inputs.input_ids, max_new_tokens=100)
+    bot_response = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    # Extract only the bot's new response (to avoid repeating context)
+    bot_response = bot_response[len(context):].strip()
+
+    # Update the conversation history
+    history.append((user_input, bot_response))
+
+    return history
+
+def gradio_format(history):
+    """
+    Format the history for Gradio ChatInterface.
+
+    Parameters:
+        history (list of tuples): Conversation history as a list of (user, bot) pairs.
+
+    Returns:
+        List of dictionaries compatible with Gradio ChatInterface.
+    """
+    return [[user, bot] for user, bot in history]
+
+# Initialize empty history
+history = []
+
+def interface_function(user_input):
+    global history
+    history = chat_with_model(history, user_input)
+    return gradio_format(history)
+
+# Create Gradio interface
+chatbot = gr.ChatInterface(
+    fn=interface_function,
+    inputs=[gr.Textbox(lines=2, label="Your Input")],
+    outputs=[gr.Chatbot(label="Chat History")],
+    title="Persian Chatbot",
+    description="A chatbot that translates or responds to Persian prompts using ISANG-v1.0-8B model."
 )
 
-# Launch the interface
-iface.launch()
+# Launch the app
+if __name__ == "__main__":
+    chatbot.launch()