Halfotter/flud

1import torch
2from transformers import AutoTokenizer, AutoModelForSequenceClassification
3import numpy as np
4
5def preprocess_function(examples, tokenizer, max_length=512):
6    """
7    Preprocess text data for the steel material classification model
8    
9    Args:
10        examples: Dataset examples containing text
11        tokenizer: Tokenizer instance
12        max_length: Maximum sequence length
13        
14    Returns:
15        dict: Tokenized inputs
16    """
17    # Tokenize the texts
18    result = tokenizer(
19        examples["text"],
20        truncation=True,
21        padding="max_length",
22        max_length=max_length,
23        return_tensors="pt"
24    )
25    
26    return result
27
28def postprocess_function(predictions, id2label):
29    """
30    Postprocess model predictions
31    
32    Args:
33        predictions: Raw model predictions
34        id2label: Mapping from label IDs to label names
35        
36    Returns:
37        dict: Processed predictions with labels and probabilities
38    """
39    # Convert logits to probabilities
40    probabilities = torch.nn.functional.softmax(torch.tensor(predictions), dim=-1)
41    
42    # Get top predictions
43    top_probs, top_indices = torch.topk(probabilities, k=5, dim=1)
44    
45    results = []
46    for i in range(len(predictions)):
47        sample_results = []
48        for j in range(5):
49            label_id = top_indices[i][j].item()
50            probability = top_probs[i][j].item()
51            label = id2label[label_id]
52            
53            sample_results.append({
54                "label": label,
55                "label_id": label_id,
56                "probability": probability
57            })
58        results.append(sample_results)
59    
60    return results
61
62def validate_input(text):
63    """
64    Validate input text for classification
65    
66    Args:
67        text: Input text to validate
68        
69    Returns:
70        bool: True if valid, False otherwise
71    """
72    if not isinstance(text, str):
73        return False
74    
75    if len(text.strip()) == 0:
76        return False
77    
78    if len(text) > 1000:  # Reasonable limit for steel material descriptions
79        return False
80    
81    return True
82
83def clean_text(text):
84    """
85    Clean and normalize input text
86    
87    Args:
88        text: Raw input text
89        
90    Returns:
91        str: Cleaned text
92    """
93    # Remove extra whitespace
94    text = " ".join(text.split())
95    
96    # Normalize Korean characters (if needed)
97    # Add any specific text cleaning rules here
98    
99    return text.strip()
100
101# Example usage
102if __name__ == "__main__":
103    # Load tokenizer
104    tokenizer = AutoTokenizer.from_pretrained(".")
105    
106    # Example preprocessing
107    example_texts = [
108        "철광석을 고로에서 환원하여 선철을 제조하는 과정",
109        "천연가스를 연료로 사용하여 고로를 가열",
110        "석회석을 첨가하여 슬래그를 형성"
111    ]
112    
113    # Clean and validate texts
114    cleaned_texts = []
115    for text in example_texts:
116        if validate_input(text):
117            cleaned_text = clean_text(text)
118            cleaned_texts.append(cleaned_text)
119    
120    # Preprocess
121    examples = {"text": cleaned_texts}
122    tokenized = preprocess_function(examples, tokenizer)
123    
124    print("=== Preprocessing Example ===")
125    print(f"Input texts: {cleaned_texts}")
126    print(f"Tokenized shape: {tokenized['input_ids'].shape}")
127    print(f"Attention mask shape: {tokenized['attention_mask'].shape}")
128