Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import torch.nn.functional as F | |
| from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
| # Initialize device | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| print(f"Using device: {device}") | |
| def load_model(model_id="ChaitraSaiK/Talk_To_Me_model"): | |
| """Load model and tokenizer from Hugging Face Hub""" | |
| try: | |
| # Load tokenizer and model from hub | |
| tokenizer = GPT2Tokenizer.from_pretrained("gpt2") | |
| model = GPT2LMHeadModel.from_pretrained( | |
| model_id, | |
| trust_remote_code=True, | |
| use_auth_token=True # If your model is private | |
| ) | |
| print(f"Model loaded from Hugging Face Hub: {model_id}") | |
| model.to(device) | |
| model.eval() | |
| return model, tokenizer | |
| except Exception as e: | |
| print(f"Error loading model: {str(e)}") | |
| print("Falling back to base GPT-2 model") | |
| # Fallback to base GPT-2 if there's an error | |
| tokenizer = GPT2Tokenizer.from_pretrained("gpt2") | |
| model = GPT2LMHeadModel.from_pretrained("gpt2") | |
| model.to(device) | |
| model.eval() | |
| return model, tokenizer | |
| # Load the model from Hugging Face Hub | |
| model, tokenizer = load_model() | |
| def generate_text(prompt, max_length=100, temperature=0.7, top_k=50): | |
| try: | |
| # Encode the prompt | |
| input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device) | |
| # Generate text | |
| with torch.no_grad(): | |
| output = model.generate( | |
| input_ids, | |
| max_length=max_length, | |
| temperature=temperature, | |
| top_k=top_k, | |
| num_return_sequences=1, | |
| pad_token_id=tokenizer.eos_token_id, | |
| do_sample=True, | |
| no_repeat_ngram_size=2, | |
| repetition_penalty=1.2 | |
| ) | |
| # Decode the generated text | |
| generated_text = tokenizer.decode(output[0], skip_special_tokens=True) | |
| return generated_text | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| # Create Gradio interface | |
| iface = gr.Interface( | |
| fn=generate_text, | |
| inputs=[ | |
| gr.Textbox( | |
| label="Prompt", | |
| placeholder="Enter your Shakespeare-style prompt here...", | |
| lines=3 | |
| ), | |
| gr.Slider( | |
| minimum=10, | |
| maximum=500, | |
| value=100, | |
| step=1, | |
| label="Max Length" | |
| ), | |
| gr.Slider( | |
| minimum=0.1, | |
| maximum=2.0, | |
| value=0.7, | |
| step=0.1, | |
| label="Temperature" | |
| ), | |
| gr.Slider( | |
| minimum=1, | |
| maximum=100, | |
| value=50, | |
| step=1, | |
| label="Top-k" | |
| ) | |
| ], | |
| outputs=gr.Textbox(label="Generated Shakespeare-style Text", lines=5), | |
| title="Shakespeare-style Text Generator", | |
| description=""" | |
| Enter a prompt and the model will generate Shakespeare-style text based on it. | |
| - Temperature: Higher values make the output more random, lower values make it more focused | |
| - Top-k: Number of highest probability tokens to consider for each generation step | |
| """ | |
| ) | |
| if __name__ == "__main__": | |
| iface.launch() |