Update app.py

app.py

@@ -1,28 +1,15 @@
 import gradio as gr
-import subprocess
-import sys
 import torch
 import os
-from gradio_client import Client
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Clone the transformers repository and install from source
-subprocess.check_call([sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps", "--upgrade", "transformers", "torch"])
-
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
-
-# Initialize ZeroGPU
-os.environ["ZEROGPU"] = "1"
-
-# Set the device to GPU if available, otherwise fallback to ZeroGPU
+# Set the device to GPU if available, otherwise use CPU
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Register the olmoe model type
-AutoConfig.register("olmoe", AutoConfig)
-AutoModelForCausalLM.register("olmoe", AutoModelForCausalLM)
-
 # Load the model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924-Instruct").to(DEVICE)
-tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924-Instruct")
+model_name = "allenai/OLMoE-1B-7B-0924-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(DEVICE)
 
 # Define the system prompt
 system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
@@ -31,21 +18,36 @@ system_prompt = ("Adopt the persona of hilariously pissed off Andrej Karpathy "
                  "without using any analogies and always showing full working code or output in his answers.")
 
 # Define a function for generating text
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
-    out = model.generate(**inputs, max_length=64)
-    return tokenizer.decode(out[0])
-
-# Function to set client for session
-def set_client_for_session(request: gr.Request):
-    x_ip_token = request.headers['x-ip-token']
-    return Client("gradio/text-to-image", headers={"X-IP-Token": x_ip_token})
+def generate_text(prompt, history):
+    full_prompt = f"{system_prompt}\n\nHuman: {prompt}\n\nAssistant:"
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(DEVICE)
+    
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=4000, do_sample=True, temperature=0.5)
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    assistant_response = response.split("Assistant:")[-1].strip()
+    
+    return assistant_response
 
 # Set up the Gradio chat interface
 with gr.Blocks() as demo:
-    client = gr.State()
-    iface = gr.ChatInterface(fn=generate_text, system_prompt=system_prompt)
-    
-    demo.load(set_client_for_session, None, client)
-    iface.launch()
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    clear = gr.Button("Clear")
+
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history):
+        user_message = history[-1][0]
+        bot_message = generate_text(user_message, history)
+        history[-1][1] = bot_message
+        return history
+
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, chatbot, chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+demo.launch(share=True)