Spaces:
Running
Running
| import gradio as gr | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| import gc | |
| # Global model and tokenizer | |
| model = None | |
| tokenizer = None | |
| def load_model(): | |
| global model, tokenizer | |
| model_name = "Spestly/Athena-1-1.5B" | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_name, | |
| low_cpu_mem_usage=True, | |
| torch_dtype=torch.float32, | |
| device_map="cpu" | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model.eval() | |
| def generate_response(input_text, history): | |
| global model, tokenizer | |
| # Load model if not loaded | |
| if model is None or tokenizer is None: | |
| load_model() | |
| try: | |
| instruction = "You are an AI assistant. Respond to the following instruction clearly and concisely:\n" | |
| formatted_input = instruction + input_text | |
| inputs = tokenizer( | |
| formatted_input, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=256 | |
| ) | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=100, | |
| do_sample=True, | |
| top_k=50, | |
| top_p=0.9, | |
| temperature=0.7, | |
| pad_token_id=tokenizer.pad_token_id, | |
| eos_token_id=tokenizer.eos_token_id, | |
| use_cache=True, | |
| repetition_penalty=1.2, | |
| num_beams=1 | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| # Manual garbage collection | |
| gc.collect() | |
| # Update history | |
| history = history + f"\nUser: {input_text}\nAthena: {response}\n" | |
| return response, history | |
| except Exception as e: | |
| return f"Error: {str(e)}", history | |
| print("Starting app...") | |
| with gr.Blocks() as demo: | |
| gr.Markdown(""" | |
| # Athena-1.1.5B Chat | |
| Smaller and Smarter. | |
| """) | |
| # Initialize state and components | |
| state = gr.State("") | |
| user_input = gr.Textbox( | |
| label="Your Message", | |
| placeholder="Type your question here...", | |
| lines=3 | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| chat_history = gr.Textbox( | |
| label="Chat History", | |
| placeholder="Chat history will appear here...", | |
| lines=10, | |
| interactive=False | |
| ) | |
| with gr.Column(scale=1): | |
| gr.Markdown("### Examples") | |
| examples = [ | |
| "What's been the role of music in human societies?", | |
| "Escribe un poema corto sobre la historia del Mediterráneo.", | |
| "Jane has 8 apples. Solve this in Python." | |
| ] | |
| example_buttons = [] | |
| for example in examples: | |
| example_button = gr.Button(example) | |
| example_button.click( | |
| lambda e=example: (e, state.value), | |
| outputs=[user_input, state] | |
| ) | |
| example_buttons.append(example_button) | |
| output_text = gr.Textbox( | |
| label="Athena's Response", | |
| lines=5, | |
| interactive=False | |
| ) | |
| submit_button = gr.Button("Submit") | |
| # Handle submission | |
| submit_button.click( | |
| generate_response, | |
| inputs=[user_input, state], | |
| outputs=[output_text, state] | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |