Update app.py

app.py

@@ -2,6 +2,7 @@
 import gradio as gr
 import transformers
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
+import torch
 import os
 
 HF_TOKEN = os.environ.get("HF_TOKEN")
@@ -12,12 +13,33 @@ model = GPT2LMHeadModel.from_pretrained('skylersterling/TopicGPT', use_auth_toke
 model.eval()
 
 # Define the function that generates text from a prompt
-def generate_text(prompt):
-  input_ids = tokenizer.encode(prompt, return_tensors='pt')
-  output = model.generate(input_ids, max_new_tokens=80, do_sample=True)
-  text = tokenizer.decode(output[0], skip_special_tokens=True)
-  return text
+def generate_text(prompt, temperature, top_p):
+    input_ids = tokenizer.encode(prompt, return_tensors='pt')
+    generated_text = prompt
+    model.eval()
+    
+    with torch.no_grad():
+        for _ in range(80):  # Generate up to 80 tokens
+            outputs = model(input_ids)
+            next_token_logits = outputs.logits[:, -1, :] / temperature
+            filtered_logits = transformers.TopPLogitsWarper(top_p)(input_ids, next_token_logits)
+            next_token = torch.argmax(filtered_logits, dim=-1).unsqueeze(0)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            generated_text += tokenizer.decode(next_token[0], skip_special_tokens=True)
+            if tokenizer.decode(next_token[0], skip_special_tokens=True).strip() == tokenizer.eos_token:
+                break
+
+    return generated_text
+
+# Create a gradio interface with a text input, sliders for temperature and top-p, and a text output
+interface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.inputs.Textbox(lines=2, placeholder='Enter your prompt here...'),
+        gr.inputs.Slider(minimum=0.1, maximum=1.0, default=1.0, label='Temperature'),
+        gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.9, label='Top-p')
+    ],
+    outputs='text'
+)
 
-# Create a gradio interface with a text input and a text output
-interface = gr.Interface(fn=generate_text, inputs='text', outputs='text')
 interface.launch()