f

app.py

@@ -1,29 +1,36 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import gradio as gr
 
-model_path = "cheberle/autotrain-35swc-b4r9z"
+# Define the model paths
+base_model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
+fine_tuned_model_name = "cheberle/autotrain-35swc-b4r9z"
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_path)
+# Load the tokenizer
+tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_name)
+
+# Load the model
 model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    device_map="auto",  # Auto-distributes model across available devices
-    torch_dtype="auto"  # Automatically selects the appropriate data type
-).eval()
+    fine_tuned_model_name,
+    device_map="auto",  # Auto-distributes model across devices
+    torch_dtype="auto", # Matches model precision
+)
 
-# Prompt content
-messages = [{"role": "user", "content": "hi"}]
+# Define the chat function
+def chat(input_text):
+    input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")  # Move input to GPU
+    output = model.generate(input_ids, max_length=100)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
-# Prepare input for the model
-input_ids = tokenizer.apply_chat_template(
-    conversation=messages, 
-    tokenize=True, 
-    add_generation_prompt=True, 
-    return_tensors='pt'
+# Create a Gradio interface
+interface = gr.Interface(
+    fn=chat,
+    inputs=gr.Textbox(lines=2, placeholder="Type your input here..."),
+    outputs="text",
+    title="Chat with DeepSeek-AutoTrain Model",
+    description="Fine-tuned version of DeepSeek-R1-Distill-Qwen-7B. Ask me anything!",
 )
 
-# Generate response
-output_ids = model.generate(input_ids.to('cuda'))  # Ensure the model uses the GPU if available
-response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
-
-# Print response
-print(response)
+# Launch the interface
+if __name__ == "__main__":
+    interface.launch()

requirements.txt

@@ -1,4 +1,5 @@
 huggingface_hub==0.25.2
 transformers
-torch
-accelerate
+accelerate
+gradio
+torch