| import gradio as gr
|
| import torch
|
| import os
|
| from model_utils import load_model, generate_text, GPTConfig
|
|
|
|
|
| try:
|
| model_path = "best_model.pt"
|
| if not os.path.exists(model_path):
|
| raise FileNotFoundError(f"Model file {model_path} not found")
|
|
|
| model = load_model(model_path)
|
| if model is None:
|
| raise ValueError("Model failed to load")
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| print(f"Using device: {device}")
|
| model = model.to(device)
|
| except Exception as e:
|
| print(f"Error loading model: {e}")
|
| model = None
|
|
|
| def predict(prompt, max_tokens, temperature, top_k):
|
| """Wrapper function for Gradio interface"""
|
| if not prompt:
|
| return "Error: Please enter a prompt"
|
| if model is None:
|
| return "Error: Model failed to load. Please check the logs."
|
|
|
| try:
|
| generated_text = generate_text(
|
| model=model,
|
| prompt=prompt.strip(),
|
| max_new_tokens=int(max_tokens),
|
| temperature=float(temperature),
|
| top_k=int(top_k)
|
| )
|
| return generated_text
|
| except Exception as e:
|
| return f"Error during generation: {str(e)}"
|
|
|
|
|
| demo = gr.Interface(
|
| fn=predict,
|
| inputs=[
|
| gr.Textbox(label="Enter your prompt", lines=3, placeholder="Type your text here..."),
|
| gr.Slider(minimum=1, maximum=200, value=50, step=1, label="Max Tokens"),
|
| gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature"),
|
| gr.Slider(minimum=1, maximum=100, value=40, step=1, label="Top-k"),
|
| ],
|
| outputs=gr.Textbox(label="Generated Text", lines=5),
|
| title="GPT Text Generation",
|
| description="Enter a text prompt and the model will generate a continuation.",
|
| examples=[
|
| ["The quick brown fox", 50, 0.8, 40],
|
| ["Once upon a time", 100, 0.9, 50],
|
| ["In the distant future", 75, 0.7, 30],
|
| ],
|
| )
|
|
|
| if __name__ == "__main__":
|
| demo.launch(share=False)
|
| else:
|
| app = demo.launch(share=False) |
