trying with just base model

app.py

@@ -1,124 +1,173 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaConfig
-import gradio as gr
+# import torch
+# from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaForCausalLM, LlamaConfig
+# import gradio as gr
 
 
-# Model IDs from Hugging Face Hub
-base_model_id = "HuggingFaceTB/SmolLM2-135M"
-instruct_model_id = "MaxBlumenfeld/smollm2-135m-bootleg-instruct-01"
+# # Model IDs from Hugging Face Hub
+# base_model_id = "HuggingFaceTB/SmolLM2-135M"
+# instruct_model_id = "MaxBlumenfeld/smollm2-135m-bootleg-instruct-01"
 
-# Load tokenizer
-base_tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+# # Load tokenizer
+# base_tokenizer = AutoTokenizer.from_pretrained(base_model_id)
 
-# Load models with explicit LLaMA architecture
-base_model = LlamaForCausalLM.from_pretrained(base_model_id)
-instruct_model = LlamaForCausalLM.from_pretrained(instruct_model_id)
+# # Load models with explicit LLaMA architecture
+# base_model = LlamaForCausalLM.from_pretrained(base_model_id)
+# instruct_model = LlamaForCausalLM.from_pretrained(instruct_model_id)
 
-def generate_response(model, tokenizer, message, temperature=0.5, max_length=200, system_prompt="", is_instruct=False):
-    # Prepare input based on model type
-    if is_instruct:
-        if system_prompt:
-            full_prompt = f"{system_prompt}\n\nHuman: {message}\nAssistant:"
-        else:
-            full_prompt = f"Human: {message}\nAssistant:"
-    else:
-        # For base model, use simpler prompt format
-        full_prompt = message
+# def generate_response(model, tokenizer, message, temperature=0.5, max_length=200, system_prompt="", is_instruct=False):
+#     # Prepare input based on model type
+#     if is_instruct:
+#         if system_prompt:
+#             full_prompt = f"{system_prompt}\n\nHuman: {message}\nAssistant:"
+#         else:
+#             full_prompt = f"Human: {message}\nAssistant:"
+#     else:
+#         # For base model, use simpler prompt format
+#         full_prompt = message
     
-    inputs = tokenizer(full_prompt, return_tensors="pt")
+#     inputs = tokenizer(full_prompt, return_tensors="pt")
     
-    with torch.no_grad():
-        outputs = model.generate(
-            inputs.input_ids,
-            max_length=max_length,
-            do_sample=True,
-            temperature=temperature,
-            top_k=50,
-            top_p=0.95,
-            num_return_sequences=1,
-            pad_token_id=tokenizer.eos_token_id  # Add padding token
-        )
+#     with torch.no_grad():
+#         outputs = model.generate(
+#             inputs.input_ids,
+#             max_length=max_length,
+#             do_sample=True,
+#             temperature=temperature,
+#             top_k=50,
+#             top_p=0.95,
+#             num_return_sequences=1,
+#             pad_token_id=tokenizer.eos_token_id  # Add padding token
+#         )
         
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+#     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    if is_instruct:
-        try:
-            response = response.split("Assistant:")[-1].strip()
-        except:
-            pass
-    else:
-        response = response[len(full_prompt):].strip()
+#     if is_instruct:
+#         try:
+#             response = response.split("Assistant:")[-1].strip()
+#         except:
+#             pass
+#     else:
+#         response = response[len(full_prompt):].strip()
         
-    return response
+#     return response
 
-def chat(message, temperature, max_length, system_prompt):
-    # Generate responses from both models
-    base_response = generate_response(
-        base_model, 
-        base_tokenizer, 
-        message, 
-        temperature, 
-        max_length, 
-        system_prompt,
-        is_instruct=False
-    )
+# def chat(message, temperature, max_length, system_prompt):
+#     # Generate responses from both models
+#     base_response = generate_response(
+#         base_model, 
+#         base_tokenizer, 
+#         message, 
+#         temperature, 
+#         max_length, 
+#         system_prompt,
+#         is_instruct=False
+#     )
     
-    instruct_response = generate_response(
-        instruct_model, 
-        base_tokenizer, 
-        message, 
-        temperature, 
-        max_length, 
-        system_prompt,
-        is_instruct=True
-    )
+#     instruct_response = generate_response(
+#         instruct_model, 
+#         base_tokenizer, 
+#         message, 
+#         temperature, 
+#         max_length, 
+#         system_prompt,
+#         is_instruct=True
+#     )
     
-    return base_response, instruct_response
+#     return base_response, instruct_response
 
-# Create Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# SmolLM2-135M Comparison Demo")
-    gr.Markdown("Compare responses between base and fine-tuned versions of SmolLM2-135M")
+# # Create Gradio interface
+# with gr.Blocks() as demo:
+#     gr.Markdown("# SmolLM2-135M Comparison Demo")
+#     gr.Markdown("Compare responses between base and fine-tuned versions of SmolLM2-135M")
     
-    with gr.Row():
-        with gr.Column():
-            message_input = gr.Textbox(label="Input Message")
-            system_prompt = gr.Textbox(
-                label="System Prompt (Optional)",
-                placeholder="Set context or personality for the model",
-                lines=3
-            )
+#     with gr.Row():
+#         with gr.Column():
+#             message_input = gr.Textbox(label="Input Message")
+#             system_prompt = gr.Textbox(
+#                 label="System Prompt (Optional)",
+#                 placeholder="Set context or personality for the model",
+#                 lines=3
+#             )
             
-        with gr.Column():
-            temperature = gr.Slider(
-                minimum=0.1, 
-                maximum=2.0, 
-                value=0.5, 
-                label="Temperature"
-            )
-            max_length = gr.Slider(
-                minimum=50, 
-                maximum=500, 
-                value=200, 
-                step=10, 
-                label="Max Length"
-            )
+#         with gr.Column():
+#             temperature = gr.Slider(
+#                 minimum=0.1, 
+#                 maximum=2.0, 
+#                 value=0.5, 
+#                 label="Temperature"
+#             )
+#             max_length = gr.Slider(
+#                 minimum=50, 
+#                 maximum=500, 
+#                 value=200, 
+#                 step=10, 
+#                 label="Max Length"
+#             )
     
-    with gr.Row():
-        with gr.Column():
-            gr.Markdown("### Base Model Response")
-            base_output = gr.Textbox(label="Base Model (SmolLM2-135M)", lines=5)
+#     with gr.Row():
+#         with gr.Column():
+#             gr.Markdown("### Base Model Response")
+#             base_output = gr.Textbox(label="Base Model (SmolLM2-135M)", lines=5)
             
-        with gr.Column():
-            gr.Markdown("### Bootleg Instruct Model Response") 
-            instruct_output = gr.Textbox(label="Fine-tuned Model", lines=5)
+#         with gr.Column():
+#             gr.Markdown("### Bootleg Instruct Model Response") 
+#             instruct_output = gr.Textbox(label="Fine-tuned Model", lines=5)
     
-    submit_btn = gr.Button("Generate Responses")
-    submit_btn.click(
-        fn=chat,
-        inputs=[message_input, temperature, max_length, system_prompt],
-        outputs=[base_output, instruct_output]
-    )
+#     submit_btn = gr.Button("Generate Responses")
+#     submit_btn.click(
+#         fn=chat,
+#         inputs=[message_input, temperature, max_length, system_prompt],
+#         outputs=[base_output, instruct_output]
+#     )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import gradio as gr
+
+model_id = "MaxBlumenfeld/smollm2-135m"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+
+def generate_response(message, temperature=0.7, max_length=200):
+   prompt = f"Human: {message}\nAssistant:"
+   inputs = tokenizer(prompt, return_tensors="pt")
+   
+   with torch.no_grad():
+       outputs = model.generate(
+           inputs.input_ids,
+           max_length=max_length,
+           temperature=temperature,
+           do_sample=True,
+           pad_token_id=tokenizer.eos_token_id
+       )
+   
+   response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+   return response.split("Assistant:")[-1].strip()
+
+with gr.Blocks() as demo:
+   gr.Markdown("# SmolLM2 Bootleg Instruct Chat")
+   
+   with gr.Row():
+       with gr.Column():
+           message = gr.Textbox(label="Message")
+           temp = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature")
+           max_len = gr.Slider(minimum=50, maximum=500, value=200, label="Max Length") 
+           submit = gr.Button("Send")
+           
+       with gr.Column():
+           output = gr.Textbox(label="Response")
+           
+   submit.click(
+       generate_response,
+       inputs=[message, temp, max_len],
+       outputs=output
+   )
 
 if __name__ == "__main__":
-    demo.launch()
+   demo.launch()
+