Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, GPT2Tokenizer | |
| import torch | |
| from huggingface_hub import login | |
| import os | |
| # Load text generation model with fallback for tokenizer | |
| def load_model(model_name): | |
| try: | |
| # Try loading the fast tokenizer first | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| except Exception as e: | |
| print(f"Fast tokenizer not available for {model_name}. Falling back to regular tokenizer. Error: {e}") | |
| # If fast tokenizer is not available, fall back to the regular tokenizer | |
| tokenizer = GPT2Tokenizer.from_pretrained(model_name) | |
| model = AutoModelForCausalLM.from_pretrained(model_name) | |
| return tokenizer, model | |
| # Load Hugging Face token | |
| hf_token = os.getenv('HF_API_TOKEN') | |
| if not hf_token: | |
| raise ValueError("Error: Hugging Face token not found. Please set it as an environment variable.") | |
| # Login to Hugging Face Hub | |
| login(hf_token) | |
| # Function to compare text generation from both models | |
| def compare_models(prompt, original_model_name, fine_tuned_model_name): | |
| # Load the original and fine-tuned models based on user input | |
| original_tokenizer, original_model = load_model(original_model_name) | |
| fine_tuned_tokenizer, fine_tuned_model = load_model(fine_tuned_model_name) | |
| # Ensure models are in evaluation mode | |
| original_model.eval() | |
| fine_tuned_model.eval() | |
| # Generate text with the original model | |
| inputs_orig = original_tokenizer(prompt, return_tensors="pt") | |
| with torch.no_grad(): | |
| generated_ids_orig = original_model.generate(inputs_orig["input_ids"], max_length=100) | |
| generated_text_orig = original_tokenizer.decode(generated_ids_orig[0], skip_special_tokens=True) | |
| # Generate text with the fine-tuned model | |
| inputs_fine = fine_tuned_tokenizer(prompt, return_tensors="pt") | |
| with torch.no_grad(): | |
| generated_ids_fine = fine_tuned_model.generate(inputs_fine["input_ids"], max_length=100) | |
| generated_text_fine = fine_tuned_tokenizer.decode(generated_ids_fine[0], skip_special_tokens=True) | |
| # Return the generated text from both models for comparison | |
| result = { | |
| "Original Model Output": generated_text_orig, | |
| "Fine-Tuned Model Output": generated_text_fine | |
| } | |
| return result | |
| # Gradio Interface | |
| iface = gr.Interface( | |
| fn=compare_models, | |
| inputs=[ | |
| gr.Textbox(lines=5, placeholder="Enter text here...", label="Input Text"), | |
| gr.Textbox(lines=1, placeholder="Enter original model name...", label="Original Model Name"), | |
| gr.Textbox(lines=1, placeholder="Enter fine-tuned model name...", label="Fine-Tuned Model Name") | |
| ], | |
| outputs=gr.JSON(label="Generated Texts"), | |
| title="Compare Text Generation from Original and Fine-Tuned Models", | |
| description="Enter a prompt and model names to generate text from the original and fine-tuned models." | |
| ) | |
| iface.launch() | |