Update app.py

app.py

@@ -1,19 +1,32 @@
 import gradio as gr
 import logging
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")
 
-def wrapped_model(*args, **kwargs):
-    logging.info(f"Input arguments: {args}")
-    logging.info(f"Keyword arguments: {kwargs}")
+model_name = "deepseek-ai/DeepSeek-R1"
+logging.info(f"Loading model: {model_name}")
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-    output = original_model(*args, **kwargs)
-    logging.info(f"Output: {output}")
-    return output
+def wrapped_model(input_text):
+    logging.info(f"Input text: {input_text}")
 
-interface = gr.load("models/deepseek-ai/DeepSeek-R1")
+    inputs = tokenizer(input_text, return_tensors="pt")
+    logging.debug(f"Tokenized input: {inputs}")
 
-original_model = interface.fn
-interface.fn = wrapped_model
+    outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)
+    logging.debug(f"Model output tokens: {outputs}")
+
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    logging.info(f"Decoded output: {decoded_output}")
+
+    return decoded_output
+
+interface = gr.Interface(
+    fn=wrapped_model,
+    inputs=gr.Textbox(lines=2, label="Enter your prompt"),
+    outputs=gr.Textbox(label="Output")
+)
 
 interface.launch()