Update app.py

app.py

@@ -1,31 +1,42 @@
+# Import the necessary libraries
 import gradio as gr
-import transformers
+import torch
 from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import os
 
+# Get the Hugging Face token from the environment variable
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 # Load the tokenizer and model
 tokenizer = GPT2Tokenizer.from_pretrained('gpt2', use_auth_token=HF_TOKEN)
 model = GPT2LMHeadModel.from_pretrained('skylersterling/TopicGPT', use_auth_token=HF_TOKEN)
 model.eval()
+model.to('cuda')
 
 # Define the function that generates text from a prompt
-def generate_text(prompt, temperature, top_p):
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
-    output = model.generate(input_ids, max_new_tokens=80, do_sample=True, temperature=temperature, top_p=top_p)
-    text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return text
+def generate_text(prompt):
+    input_tokens = tokenizer.encode(prompt, return_tensors='pt')
+    input_tokens = input_tokens.to('cuda')
+    
+    generated_tokens = []
 
-# Create a gradio interface with a text input, a text output, and sliders for temperature and top-p
-interface = gr.Interface(
-    fn=generate_text, 
-    inputs=[
-        gr.inputs.Textbox(label="Prompt"),
-        gr.inputs.Slider(minimum=0.1, maximum=1.0, step=0.1, default=1.0, label="Temperature"),
-        gr.inputs.Slider(minimum=0.1, maximum=1.0, step=0.1, default=1.0, label="Top-p")
-    ], 
-    outputs='text'
-)
+    for _ in range(80):  # Adjust the range to control the number of tokens generated
+        with torch.no_grad():
+            outputs = model(input_tokens)
+            predictions = outputs.logits
+            next_token = torch.multinomial(torch.softmax(predictions[:, -1, :], dim=-1), 1)
+        
+        generated_tokens.append(next_token.item())
+        input_tokens = torch.cat((input_tokens, next_token), dim=1)
+        
+        decoded_token = tokenizer.decode(next_token.item())
+        # Print each token as it is generated
+        print(decoded_token, end='', flush=True)
+    
+    # Decode the generated tokens to a string
+    generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True)
+    return generated_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()