Update app.py

app.py

@@ -1,48 +1,53 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-import os
-
-# Get the Hugging Face token from environment variables
-hf_token = os.getenv("API_KEY")
-
-
-
-# Load model and tokenizer
-model_name = "mistralai/Mistral-7B-v0.1"
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    use_auth_token=hf_token  
-)
-tokenizer = AutoTokenizer.from_pretrained(
-    model_name,
-    use_auth_token=hf_token  
-)
-
-# Define the generation function
-def generate_response(prompt):
-    # Tokenize input text
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-
-    # Generate response
-    generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=True)
+
+gr.load("models/mistralai/Mistral-7B-v0.1").launch()
+
+
+# import gradio as gr
+# from transformers import AutoModelForCausalLM, AutoTokenizer
+# import torch
+# import os
+
+# # Get the Hugging Face token from environment variables
+# hf_token = os.getenv("API_KEY")
+
+
+
+# # Load model and tokenizer
+# model_name = "mistralai/Mistral-7B-v0.1"
+# model = AutoModelForCausalLM.from_pretrained(
+#     model_name,
+#     device_map="auto",
+#     use_auth_token=hf_token  
+# )
+# tokenizer = AutoTokenizer.from_pretrained(
+#     model_name,
+#     use_auth_token=hf_token  
+# )
+
+# # Define the generation function
+# def generate_response(prompt):
+#     # Tokenize input text
+#     inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
+#     # Generate response
+#     generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=True)
     
-    # Decode and return response
-    return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+#     # Decode and return response
+#     return tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
-# Set up Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Text Generation")
-    input_text = gr.Textbox(placeholder="Enter your input here", lines=2)
-    output_text = gr.Textbox(label="Generated Output", lines=2)
-    submit_btn = gr.Button("Generate")
+# # Set up Gradio interface
+# with gr.Blocks() as demo:
+#     gr.Markdown("# Text Generation")
+#     input_text = gr.Textbox(placeholder="Enter your input here", lines=2)
+#     output_text = gr.Textbox(label="Generated Output", lines=2)
+#     submit_btn = gr.Button("Generate")
 
-    submit_btn.click(generate_response, inputs=input_text, outputs=output_text)
+#     submit_btn.click(generate_response, inputs=input_text, outputs=output_text)
 
-# Launch the interface
-if __name__ == "__main__":
-    demo.launch()
+# # Launch the interface
+# if __name__ == "__main__":
+#     demo.launch()