Update app.py

app.py

@@ -1,20 +1,39 @@
-import gradio as gr
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+model_id = "AlphaMindQ/GPTx-v6-60gb"
 
-tokenizer = GPT2Tokenizer.from_pretrained("AlphaMindQ/GPTx-v6-60gb")
-model = GPT2LMHeadModel.from_pretrained("AlphaMindQ/GPTx-v6-60gb")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
 model.eval()
-torch.set_num_threads(2)
 
-def generate(prompt):
+def generate_text(prompt, temperature, top_k, top_p, max_tokens, repetition_penalty):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(inputs["input_ids"], max_new_tokens=50)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
-gr.Interface(fn=generate, 
-             inputs=gr.Textbox(label="Enter prompt"), 
-             outputs=gr.Textbox(label="Generated text"),
-             title="AlphaMindQ - GPTx-v6-60gb",
-             description="Tiny GPT-2-compatible model demo"
+gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Prompt", placeholder="Type something here...", lines=4),
+        gr.Slider(0.1, 1.5, value=1.0, step=0.1, label="Temperature"),
+        gr.Slider(1, 100, value=50, step=1, label="Top-K"),
+        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P"),
+        gr.Slider(10, 512, value=128, step=1, label="Max New Tokens"),
+        gr.Slider(0.5, 2.0, value=1.0, step=0.1, label="Repetition Penalty")
+    ],
+    outputs=gr.Textbox(label="Generated Text"),
+    title="🧠 AlphaMindQ GPTx-v6-60GB",
+    theme="default"
 ).launch()