Update app.py

app.py

@@ -1,13 +1,12 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import InferenceApi
+import os
 
-# Initialize the model and tokenizer
-model_name = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+api_token = os.getenv('token')
 
-initial_scenario = "You wake up in a strange forest. Paths lead in all directions."
-history = [initial_scenario]
+client = InferenceApi("gpt2", api_token=api_token)
+
+scenario = "You wake up in a strange forest. Paths lead in all directions."
 choices = [
     "Walk north into the dense forest.",
     "Go south back towards a distant town.",
@@ -15,36 +14,31 @@ choices = [
     "Move west where the land rises."
 ]
 
-def generate_continuation(prompt):
+history = []
+
+def generate_continuation(choice):
+    global scenario, history
+    history.append(choice)
+    prompt = f"{scenario} {' '.join(history)}"
     max_length = 50  # Adjust as needed
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
-    continuation = tokenizer.decode(output[0], skip_special_tokens=True)
-    return continuation
 
-def update_scenario(choice_index):
-    global history, choices
-    chosen_action = choices[choice_index]
-    history.append(chosen_action)
-    full_story = " ".join(history)
-    new_part = generate_continuation(full_story)
-    history.append(new_part)  # Update history with the new part of the story
-    full_story = " ".join(history)
+    response = client(prompt, max_length=max_length, remove_input=True)
+    continuation = response[0]['generated_text']
 
-    # For simplicity, we keep the same choices here, but you could generate new choices based on the story.
-    return full_story
+    return continuation
 
 with gr.Blocks() as app:
     with gr.Row():
-        story_display = gr.Textbox(value=initial_scenario, show_label=False, lines=10)
+        story_display = gr.Textbox(value=scenario, lines=7)
         choice_buttons = [gr.Button(choice) for choice in choices]
     
     def handle_choice(choice_index):
-        new_story = update_scenario(choice_index)
-        story_display.update(value=new_story)
-        return [gr.update() for _ in choices]
+        chosen_action = choices[choice_index]
+        new_scenario = generate_continuation(chosen_action)
+        story_display.update(value=new_scenario)
+        return gr.update()
 
     for i, button in enumerate(choice_buttons):
-        button.click(fn=lambda x=i: handle_choice(x), outputs=story_display)
+        button.click(fn=lambda x=i: handle_choice(x), inputs=None, outputs=story_display)
 
 app.launch()